import logging
from pathlib import Path

import numpy as np
from rich.progress import Progress

from .._omegaconf_interpreter import sample_config
from ..kinetic.models import make_compartment_model

from .curves import gaussian_noise, generate_aif, generate_time
from .io import SimulationParquetWriter, batch_to_dataframe
from .measured_inputs import MeasuredAIFPool, load_measured_aif_pool


def _get_measured_pool(cfg, *, config_dir: Path | None) -> MeasuredAIFPool | None:
    aif_cfg = cfg.get("aif")
    if aif_cfg is None:
        return None

    measured_files = aif_cfg.get("measured_files")
    if not measured_files:
        return None

    entries = [str(p) for p in measured_files]
    pool = load_measured_aif_pool(entries, base_dir=config_dir)
    logging.info("Loaded %s measured AIF/TIME vectors from config.", pool.n_items)
    return pool


def _infer_synthetic_frame_count(cfg) -> int:
    protocols = cfg.time.protocols
    frame_counts = []
    for _, spec in protocols.items():
        frame_counts.append(sum(int(chunk[0]) for chunk in spec))

    unique_counts = sorted(set(frame_counts))
    if len(unique_counts) != 1:
        raise ValueError(
            "All configured protocols must have the same number of frames. "
            f"Found frame counts: {unique_counts}"
        )
    return unique_counts[0]


def run_simulations(
    cfg,
    *,
    out_dir,
    batch_size=10_000,
    config_dir: Path | None = None,
):
    rng = np.random.default_rng(cfg.seed)
    ignore = ("time.frame_durations_s",)
    total = cfg.num_simulations
    measured_pool = _get_measured_pool(cfg, config_dir=config_dir)
    n_frames = (
        measured_pool.n_frames
        if measured_pool is not None
        else _infer_synthetic_frame_count(cfg)
    )

    with Progress() as progress, SimulationParquetWriter(out_dir=out_dir) as writer:
        task = progress.add_task("[cyan]Simulating...", total=total)
        batch_idx = 0
        sample_id0 = 0
        for start in range(0, total, batch_size):
            n = min(batch_size, total - start)

            batch = {}
            batch["batch_id"] = np.full(n, batch_idx, dtype=np.int32)
            batch["sample_id"] = np.arange(sample_id0, sample_id0 + n, dtype=np.int64)
            batch["protocol"] = np.empty(n, dtype=object)
            batch["aif_source"] = np.empty(n, dtype=object)
            batch["regime"] = np.empty(n, dtype=object)
            batch["noise"] = np.empty(n, dtype=np.float32)
            batch["n_frames_used"] = np.empty(n, dtype=np.int32)
            batch["K1"] = np.empty(n, dtype=np.float32)
            batch["k2"] = np.empty(n, dtype=np.float32)
            batch["k3"] = np.empty(n, dtype=np.float32)
            batch["k4"] = np.empty(n, dtype=np.float32)
            batch["vB"] = np.empty(n, dtype=np.float32)
            batch["Tinj_s"] = np.full(n, np.nan, dtype=np.float32)
            batch["tau"] = np.full(n, np.nan, dtype=np.float32)
            batch["A1"] = np.full(n, np.nan, dtype=np.float32)
            batch["r2"] = np.full(n, np.nan, dtype=np.float32)
            batch["r3"] = np.full(n, np.nan, dtype=np.float32)
            batch["L1"] = np.full(n, np.nan, dtype=np.float32)
            batch["L2"] = np.full(n, np.nan, dtype=np.float32)
            batch["L3"] = np.full(n, np.nan, dtype=np.float32)

            batch["time"] = np.empty((n, n_frames), dtype=np.float32)
            batch["aif"] = np.empty((n, n_frames), dtype=np.float32)
            batch["tac"] = np.empty((n, n_frames), dtype=np.float32)

            for i in range(n):
                sampled = sample_config(cfg, rng, ignore_paths=ignore)

                time_params = sampled["time"]
                scan_duration = time_params["scan_duration"]
                protocol = time_params["protocol"]

                cm_cfg = sampled["compartment_model"]
                regime = cm_cfg["regime"]
                cm_params = cm_cfg["regimes"][regime]
                cm_choice = cm_cfg["num_k"]
                cm_noise = cm_cfg["noise"]

                if measured_pool is None:
                    (
                        t_fine_min,
                        t_fine_s,
                        time_min,
                        frame_durations_s,
                        dt_s,
                    ) = generate_time(
                        frame_durations_s=time_params["protocols"][protocol],
                        dt_s=time_params["dt_s"],
                    )
                    aif_params = sampled["aif"]["params"]
                    aif = generate_aif(
                        t_fine_min,
                        t_fine_s,
                        frame_durations_s,
                        dt_s,
                        **aif_params,
                    )
                    aif_source = "__simulated__"
                else:
                    aif, time_min, aif_source = measured_pool.sample(rng)
                    protocol = "measured"

                compartment_model = make_compartment_model(cm_choice)
                tac = compartment_model(aif, time_min, **cm_params)
                tac = gaussian_noise(tac, rng, factor=cm_noise)

                n_frames_used = np.sum(time_min < scan_duration)

                batch["protocol"][i] = protocol
                batch["aif_source"][i] = aif_source
                batch["regime"][i] = regime
                batch["noise"][i] = cm_noise
                batch["n_frames_used"][i] = n_frames_used
                batch["K1"][i] = cm_params["K1"]
                batch["k2"][i] = cm_params["k2"]
                batch["k3"][i] = cm_params["k3"]
                batch["k4"][i] = cm_params["k4"]
                batch["vB"][i] = cm_params["vB"]
                if measured_pool is None:
                    batch["Tinj_s"][i] = aif_params["Tinj_s"]
                    batch["tau"][i] = aif_params["tau"]
                    batch["A1"][i] = aif_params["A1"]
                    batch["r2"][i] = aif_params["r2"]
                    batch["r3"][i] = aif_params["r3"]
                    batch["L1"][i] = aif_params["L1"]
                    batch["L2"][i] = aif_params["L2"]
                    batch["L3"][i] = aif_params["L3"]

                batch["time"][i] = time_min
                batch["aif"][i] = aif
                batch["tac"][i] = tac

                progress.update(task, advance=1)

            frame = batch_to_dataframe(batch)
            writer.write_batch(frame)
            del batch

            n_batches = (total + batch_size - 1) // batch_size
            logging.debug("%s/%s batches complete...", batch_idx + 1, n_batches)

            sample_id0 += n
            batch_idx += 1
