from itertools import product
from pathlib import Path
import pytask

import numpy as np
import pandas as pd
import pyroll.basic as pr
from scipy import stats

from esaform2023.weiner.config import DATA_DIR
from esaform2023.weiner.sim.config import create_sequence, FIELDS, COUNT, IN_TEMP_MEAN, IN_TEMP_STD, IN_DIA_STD, \
    IN_DIA_MEAN, create_in_profile

SEED = 38945729345645209345
RNG = np.random.default_rng(SEED)


@pytask.mark.depends_on({"config": "config.py", "dist": DATA_DIR / "duo_pauses_dist.csv"})
@pytask.mark.produces("sim2_results.csv")
def task_sim2(produces: Path, depends_on: dict[str, Path]):
    in_temps = np.concatenate([[IN_TEMP_MEAN], stats.norm.rvs(IN_TEMP_MEAN, IN_TEMP_STD, COUNT, random_state=RNG)])
    in_dia = np.concatenate([[IN_DIA_MEAN], stats.norm.rvs(IN_DIA_MEAN, IN_DIA_STD, COUNT, random_state=RNG)])

    in_profiles = [
        create_in_profile(
            diameter=d,
            temperature=t,
        )
        for t, d in zip(in_temps, in_dia)
    ]

    sequence = create_sequence()
    roll_passes = [u for u in sequence if isinstance(u, pr.RollPass)]
    transports = [u for u in sequence if isinstance(u, pr.Transport)][0:-3]

    dist = pd.read_csv(depends_on["dist"], header=0)

    count_of_transports = len(transports)
    pauses = np.concatenate([
        [u.duration for u in sequence if isinstance(u, pr.Transport)],
        stats.gamma.rvs(a=dist["alpha"], scale=dist["beta"], loc=dist["loc"],
                        size=COUNT * count_of_transports, random_state=RNG)
    ])
    pauses.resize((COUNT + 1, count_of_transports))

    df = pd.DataFrame(
        {
            ("in", "diameter"): in_dia,
            ("in", "temperature"): in_temps,
        } | {
            ("pause", i): p
            for i, p in enumerate(pauses.T)
        } | {
            (f, p.label): np.full(COUNT + 1, np.NaN)
            for f, p in product(FIELDS, roll_passes)
        }
    )

    for i, ip in enumerate(in_profiles):
        for t, p in zip(transports, pauses[i]):
            t.duration = p

        pr.solve(sequence, ip)

        for p in roll_passes:
            for k, f in FIELDS.items():
                df[k, p.label][i] = f(p)

    df.to_csv(produces)
