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.sim.config import FIELDS, COUNT, IN_TEMP_MEAN, IN_TEMP_STD, IN_DIA_STD, IN_DIA_MEAN, \
    create_sequence, create_in_profile

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


@pytask.mark.depends_on("config.py")
@pytask.mark.produces("sim1_results.csv")
def task_sim1(produces: 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)]

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

    for i, ip in enumerate(in_profiles):
        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)
