from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping

from app.simulation.components.amesim.flow.pipes import AmesimPnl0002
from app.simulation.reporting.amesim_results import (
    AmesimResults,
    load_test_mql_amesim_results,
)


@dataclass(frozen=True, slots=True)
class Pnl0002ReplayPaths:
    """AMESim data paths needed to replay one PNL0002 resistance."""

    center_pressure: str
    center_temperature: str
    port_1_pressure: str
    port_1_temperature: str
    port_1_mass_flow: str
    port_2_pressure: str
    port_2_temperature: str
    port_2_mass_flow: str
    reynolds: str
    friction_factor: str


PNL83_REPLAY_PATHS = Pnl0002ReplayPaths(
    center_pressure="pctr@pneumatic_83",
    center_temperature="tctr@pneumatic_83",
    port_1_pressure="press1@pnnode4_16",
    port_1_temperature="temp1@pnnode4_16",
    port_1_mass_flow="dm1@pneumatic_83",
    port_2_pressure="press3@pnnode4_17",
    port_2_temperature="temp3@pnnode4_17",
    port_2_mass_flow="dm2@pneumatic_83",
    reynolds="re@pneumatic_83",
    friction_factor="ff@pneumatic_83",
)


def _required_series(
    results: AmesimResults,
    path: str,
) -> tuple[float, ...]:
    try:
        values = results.series(path)
    except KeyError as exc:
        raise ValueError(f"AMESim replay variable is not saved: {path}") from exc
    if len(values) != len(results.times):
        raise ValueError(f"AMESim replay variable has an invalid length: {path}")
    return values


def _metric_summary(
    rows: list[dict[str, float]],
    key: str,
) -> dict[str, float]:
    values = [float(row[key]) for row in rows]
    max_index = max(range(len(values)), key=lambda index: abs(values[index]))
    return {
        "maxAbs": abs(values[max_index]),
        "maxAbsTime": rows[max_index]["time"],
        "finalSigned": values[-1],
    }


def replay_pnl0002_amesim_states(
    pipe: AmesimPnl0002,
    results: AmesimResults,
    paths: Pnl0002ReplayPaths,
    *,
    amesim_mass_flow_scale: float = -1.0e-3,
) -> dict[str, object]:
    """Replay saved AMESim states through current PNL0002 flow functions.

    This is a calibration-only, no-integration calculation. It does not write
    states into the pipe or alter the production simulation path.
    """

    series_by_field = {
        field: _required_series(results, getattr(paths, field))
        for field in paths.__dataclass_fields__
    }
    rows: list[dict[str, float]] = []
    for index, time_s in enumerate(results.times):
        center_pressure = series_by_field["center_pressure"][index]
        center_temperature = series_by_field["center_temperature"][index]
        port_1_pressure = series_by_field["port_1_pressure"][index]
        port_2_pressure = series_by_field["port_2_pressure"][index]
        observed_flow_1 = (
            amesim_mass_flow_scale
            * series_by_field["port_1_mass_flow"][index]
        )
        observed_flow_2 = (
            amesim_mass_flow_scale
            * series_by_field["port_2_mass_flow"][index]
        )
        upstream_temperature_1 = (
            series_by_field["port_1_temperature"][index]
            if observed_flow_1 >= 0.0
            else center_temperature
        )
        upstream_temperature_2 = (
            series_by_field["port_2_temperature"][index]
            if observed_flow_2 >= 0.0
            else center_temperature
        )
        predicted_flow_1 = pipe.mass_flow(
            port_1_pressure,
            center_pressure,
            upstream_temperature_1,
        )
        predicted_flow_2 = pipe.mass_flow(
            port_2_pressure,
            center_pressure,
            upstream_temperature_2,
        )
        reynolds_1 = pipe.reynolds_number(
            observed_flow_1,
            upstream_temperature_1,
        )
        reynolds_2 = pipe.reynolds_number(
            observed_flow_2,
            upstream_temperature_2,
        )
        friction_1 = pipe.friction_factor(reynolds_1)
        friction_2 = pipe.friction_factor(reynolds_2)
        replay_reynolds = 0.5 * (reynolds_1 + reynolds_2)
        replay_friction = 0.5 * (friction_1 + friction_2)
        rows.append(
            {
                "time": float(time_s),
                "centerPressure": center_pressure,
                "centerTemperature": center_temperature,
                "port1Pressure": port_1_pressure,
                "port2Pressure": port_2_pressure,
                "observedPort1MassFlow": observed_flow_1,
                "observedPort2MassFlow": observed_flow_2,
                "predictedPort1MassFlow": predicted_flow_1,
                "predictedPort2MassFlow": predicted_flow_2,
                "port1MassFlowError": predicted_flow_1 - observed_flow_1,
                "port2MassFlowError": predicted_flow_2 - observed_flow_2,
                "amesimReynolds": series_by_field["reynolds"][index],
                "replayReynolds": replay_reynolds,
                "reynoldsError": (
                    replay_reynolds - series_by_field["reynolds"][index]
                ),
                "amesimFrictionFactor": series_by_field["friction_factor"][index],
                "replayFrictionFactor": replay_friction,
                "frictionFactorError": (
                    replay_friction
                    - series_by_field["friction_factor"][index]
                ),
            }
        )

    metric_keys = (
        "port1MassFlowError",
        "port2MassFlowError",
        "reynoldsError",
        "frictionFactorError",
    )
    return {
        "mode": "amesim-state-replay-no-integration",
        "component": pipe.name,
        "pointCount": len(rows),
        "massFlowScale": amesim_mass_flow_scale,
        "paths": {
            field: getattr(paths, field)
            for field in paths.__dataclass_fields__
        },
        "parameters": dict(pipe.parameter_values),
        "summary": {
            key: _metric_summary(rows, key)
            for key in metric_keys
        },
        "rows": rows,
    }


def _compile_project_pipe(
    project_path: Path,
    component_name: str,
) -> AmesimPnl0002:
    from app.main import (
        ReactFlowProjectPayload,
        _compile_xml_document_or_422,
        _validated_xml_document_or_422,
        build_reactflow_system_xml,
        validate_system_xml_document,
    )

    payload = ReactFlowProjectPayload.model_validate_json(
        project_path.read_text(encoding="utf-8")
    )
    xml_bytes = build_reactflow_system_xml(payload)
    document = _validated_xml_document_or_422(
        validate_system_xml_document(xml_bytes)
    )
    _project, network = _compile_xml_document_or_422(document)
    try:
        component = network.components[component_name]
    except KeyError as exc:
        raise ValueError(f"Project component does not exist: {component_name}") from exc
    if not isinstance(component, AmesimPnl0002):
        raise ValueError(f"Project component is not PNL0002: {component_name}")
    return component


def _paths_from_arguments(arguments: argparse.Namespace) -> Pnl0002ReplayPaths:
    values: Mapping[str, str] = {
        field: getattr(arguments, field)
        for field in PNL83_REPLAY_PATHS.__dataclass_fields__
    }
    return Pnl0002ReplayPaths(**values)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Replay saved AMESim p/T/m_flow through a PNL0002 model without integration."
    )
    parser.add_argument("project", type=Path)
    parser.add_argument("amesim_archive", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--component", default="pneumatic_83")
    parser.add_argument("--mass-flow-scale", type=float, default=-1.0e-3)
    for field in PNL83_REPLAY_PATHS.__dataclass_fields__:
        parser.add_argument(
            "--" + field.replace("_", "-"),
            dest=field,
            default=getattr(PNL83_REPLAY_PATHS, field),
        )
    arguments = parser.parse_args()

    pipe = _compile_project_pipe(arguments.project, arguments.component)
    results = load_test_mql_amesim_results(arguments.amesim_archive)
    report = replay_pnl0002_amesim_states(
        pipe,
        results,
        _paths_from_arguments(arguments),
        amesim_mass_flow_scale=arguments.mass_flow_scale,
    )
    arguments.output.parent.mkdir(parents=True, exist_ok=True)
    arguments.output.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({key: report[key] for key in (
        "mode",
        "component",
        "pointCount",
        "parameters",
        "summary",
    )}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()