from __future__ import annotations

from bisect import bisect_left
import csv
from dataclasses import dataclass
from pathlib import Path
from typing import Any


PRIMARY_KEYS = (
    "mytank.p",
    "mytank.T",
    "mycylinder.p",
    "mycylinder.T",
)

MODELICA_COMPARISON_COLUMNS = {
    "mytank.p": "mytank.p",
    "mytank.T": "mytank.T",
    "mycylinder.p": "mycylinder.p",
    "mycylinder.T": "mycylinder.T",
    "branch.upper_branch.p": "mypipe.p",
    "branch.upper_branch.in": "myorifice.port_a.m_flow",
    "branch.upper_branch.out": "mytee1.port_out2.m_flow",
    "branch.lower_branch.p": "mypipe1.p",
    "branch.lower_branch.in": "myorifice1.port_a.m_flow",
    "branch.lower_branch.out": "mytee1.port_out1.m_flow",
}

COMPARISON_KEYS = tuple(MODELICA_COMPARISON_COLUMNS.keys())


def _branch_series_values(
    series: dict[str, list[float]],
    branch_name: str,
    legacy_key: str,
) -> list[float]:
    generic_key = f"branch.{branch_name}.{legacy_key.split('.')[-1]}"
    if generic_key in series:
        return series[generic_key]
    return series[legacy_key]


@dataclass(frozen=True)
class TestModelArtifacts:
    primary_csv_path: Path
    temperature_csv_path: Path
    temperature_svg_path: Path
    run_report_path: Path
    comparison_csv_path: Path | None = None
    comparison_summary_path: Path | None = None


def format_testmodel_run_report(
    *,
    network_summary: str,
    initialization: Any,
    raw_initial_state: tuple[float, ...],
    consistent_initial_state: tuple[float, ...],
    solution: Any,
    series: dict[str, list[float]],
    solve_diagnostics: Any,
    artifacts: TestModelArtifacts,
    comparison_summary: dict[str, tuple[float, float]] | None,
) -> str:
    lines = [
        network_summary,
        "",
        f"Initialization converged: {initialization.converged}",
        f"Initialization iterations: {initialization.iterations}",
        f"Initialization max state delta: {initialization.max_state_delta:.6e}",
        f"Initialization max flow delta: {initialization.max_flow_delta:.6e}",
        f"Initialization max enthalpy delta: {initialization.max_enthalpy_delta:.6e}",
        (
            "Initialization downstream pressure spread: "
            f"{initialization.downstream_pressure_spread:.6e}"
        ),
        "",
        "Raw initial state vector:",
        str(list(raw_initial_state)),
        "",
        "Constraint-consistent initial state vector:",
        str(list(consistent_initial_state)),
        "",
        f"Solver success: {solution.success}",
        f"Solver message: {solution.message}",
        f"Final time: {solution.t[-1]:.2f} s",
        f"Final tank pressure: {series['mytank.p'][-1]:.3f} Pa",
        f"Final tank temperature: {series['mytank.T'][-1]:.3f} K",
        f"Final cylinder pressure: {series['mycylinder.p'][-1]:.3f} Pa",
        (
            "Final branch inflow: "
            f"{_branch_series_values(series, 'upper_branch', 'branch_upper.in')[-1] + _branch_series_values(series, 'lower_branch', 'branch_lower.in')[-1]:.6f} kg/s"
        ),
    ]

    if solve_diagnostics is not None:
        lines.extend(
            [
                "",
                "Final closure solve diagnostics:",
                (
                    "Upper branch inlet solve: "
                    f"converged={solve_diagnostics.upper_branch_inlet.converged}, "
                    f"iterations={solve_diagnostics.upper_branch_inlet.iterations}, "
                    f"residual={solve_diagnostics.upper_branch_inlet.residual:.6e}"
                ),
                (
                    "Lower branch inlet solve: "
                    f"converged={solve_diagnostics.lower_branch_inlet.converged}, "
                    f"iterations={solve_diagnostics.lower_branch_inlet.iterations}, "
                    f"residual={solve_diagnostics.lower_branch_inlet.residual:.6e}"
                ),
            ]
        )
        if solve_diagnostics.downstream_pressure_projection is not None:
            lines.append(
                "Downstream pressure projection: "
                f"converged={solve_diagnostics.downstream_pressure_projection.converged}, "
                f"iterations={solve_diagnostics.downstream_pressure_projection.iterations}, "
                f"residual={solve_diagnostics.downstream_pressure_projection.residual:.6e}"
            )

    lines.extend(
        [
        f"Primary series CSV: {artifacts.primary_csv_path}",
        f"Temperature CSV: {artifacts.temperature_csv_path}",
        f"Temperature plot: {artifacts.temperature_svg_path}",
        f"Run report TXT: {artifacts.run_report_path}",
        ]
    )

    if (
        artifacts.comparison_csv_path is not None
        and artifacts.comparison_summary_path is not None
    ):
        lines.extend(
            [
                f"Modelica comparison CSV: {artifacts.comparison_csv_path}",
                f"Modelica comparison summary: {artifacts.comparison_summary_path}",
            ]
        )
        if comparison_summary is not None:
            for key, (max_abs_error, max_rel_error) in comparison_summary.items():
                lines.append(
                    f"{key} max abs error: {max_abs_error:.6f}, "
                    f"max rel error: {max_rel_error:.6%}"
                )

    return "\n".join(lines) + "\n"


def write_testmodel_run_report(output_dir: Path, report_text: str) -> Path:
    report_path = output_dir / "testmodel_run_report.txt"
    report_path.write_text(report_text, encoding="utf-8")
    return report_path


def _write_primary_series_csv(output_dir: Path, series: dict[str, list[float]]) -> Path:
    csv_path = output_dir / "testmodel_primary_series.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["time_s", *PRIMARY_KEYS])
        for index, time_value in enumerate(series["time"]):
            writer.writerow([time_value, *(series[key][index] for key in PRIMARY_KEYS)])
    return csv_path


def _write_temperature_csv(output_dir: Path, time_values: list[float], temperatures: list[float]) -> Path:
    csv_path = output_dir / "testmodel_tank_temperature.csv"
    with csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["time_s", "mytank_T_K"])
        writer.writerows(zip(time_values, temperatures))
    return csv_path


def _write_temperature_svg(output_dir: Path, time_values: list[float], temperatures: list[float]) -> Path:
    svg_path = output_dir / "testmodel_tank_temperature.svg"

    width = 900
    height = 520
    left = 90
    right = 40
    top = 60
    bottom = 70
    plot_width = width - left - right
    plot_height = height - top - bottom

    min_time = min(time_values)
    max_time = max(time_values)
    min_temp = min(temperatures)
    max_temp = max(temperatures)
    temp_padding = max(1.0, (max_temp - min_temp) * 0.08)
    min_temp -= temp_padding
    max_temp += temp_padding

    def scale_x(value: float) -> float:
        return left + (value - min_time) / max(max_time - min_time, 1e-12) * plot_width

    def scale_y(value: float) -> float:
        return top + (max_temp - value) / max(max_temp - min_temp, 1e-12) * plot_height

    points = " ".join(
        f"{scale_x(time_value):.2f},{scale_y(temperature):.2f}"
        for time_value, temperature in zip(time_values, temperatures)
    )

    x_ticks = 5
    y_ticks = 5
    x_tick_markup = []
    y_tick_markup = []

    for index in range(x_ticks + 1):
        fraction = index / x_ticks
        time_value = min_time + fraction * (max_time - min_time)
        x = left + fraction * plot_width
        x_tick_markup.append(
            f'<line x1="{x:.2f}" y1="{top}" x2="{x:.2f}" y2="{top + plot_height}" '
            'stroke="#d9e2ec" stroke-width="1" />'
        )
        x_tick_markup.append(
            f'<text x="{x:.2f}" y="{height - 30}" text-anchor="middle" '
            'font-size="14" fill="#102a43">'
            f"{time_value:.1f}</text>"
        )

    for index in range(y_ticks + 1):
        fraction = index / y_ticks
        temp_value = min_temp + fraction * (max_temp - min_temp)
        y = top + plot_height - fraction * plot_height
        y_tick_markup.append(
            f'<line x1="{left}" y1="{y:.2f}" x2="{left + plot_width}" y2="{y:.2f}" '
            'stroke="#d9e2ec" stroke-width="1" />'
        )
        y_tick_markup.append(
            f'<text x="{left - 12}" y="{y + 5:.2f}" text-anchor="end" '
            'font-size="14" fill="#102a43">'
            f"{temp_value:.1f}</text>"
        )

    svg_content = f"""<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">
  <rect width="{width}" height="{height}" fill="#f7fafc" rx="18" ry="18" />
  <text x="{width / 2:.0f}" y="32" text-anchor="middle" font-size="24" fill="#102a43">Python Testmodel Tank Temperature</text>
  <text x="{width / 2:.0f}" y="{height - 8}" text-anchor="middle" font-size="16" fill="#486581">Time (s)</text>
  <text x="26" y="{height / 2:.0f}" text-anchor="middle" font-size="16" fill="#486581" transform="rotate(-90 26 {height / 2:.0f})">Temperature (K)</text>
  <rect x="{left}" y="{top}" width="{plot_width}" height="{plot_height}" fill="#ffffff" stroke="#bcccdc" stroke-width="1.5" />
  {''.join(x_tick_markup)}
  {''.join(y_tick_markup)}
  <polyline fill="none" stroke="#d64545" stroke-width="3" stroke-linejoin="round" stroke-linecap="round" points="{points}" />
</svg>
"""
    svg_path.write_text(svg_content, encoding="utf-8")
    return svg_path


def load_modelica_series(csv_path: Path, variable_names: tuple[str, ...]) -> dict[str, list[float]]:
    series = {"time": []}
    for variable_name in variable_names:
        series[variable_name] = []

    with csv_path.open("r", newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        available_variable_names = tuple(
            variable_name
            for variable_name in variable_names
            if MODELICA_COMPARISON_COLUMNS.get(variable_name, variable_name) in (reader.fieldnames or ())
        )
        for row in reader:
            series["time"].append(float(row["time"]))
            for variable_name in available_variable_names:
                modelica_column = MODELICA_COMPARISON_COLUMNS.get(variable_name, variable_name)
                series[variable_name].append(float(row[modelica_column]))

    return series


def _interpolate_series_value(time_values: list[float], values: list[float], target_time: float) -> float:
    if target_time <= time_values[0]:
        return values[0]
    if target_time >= time_values[-1]:
        return values[-1]

    right_index = bisect_left(time_values, target_time)
    if right_index < len(time_values) and abs(time_values[right_index] - target_time) <= 1e-12:
        return values[right_index]

    left_index = right_index - 1
    left_time = time_values[left_index]
    right_time = time_values[right_index]
    fraction = (target_time - left_time) / (right_time - left_time)
    return values[left_index] + fraction * (values[right_index] - values[left_index])


def write_modelica_comparison(
    output_dir: Path,
    python_series: dict[str, list[float]],
    modelica_series: dict[str, list[float]],
) -> tuple[Path, Path, dict[str, tuple[float, float]]]:
    comparison_csv_path = output_dir / "testmodel_modelica_comparison.csv"
    summary_path = output_dir / "testmodel_modelica_comparison_summary.txt"
    summary: dict[str, tuple[float, float]] = {}

    with comparison_csv_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        header = ["time_s"]
        comparison_keys = tuple(
            key
            for key in COMPARISON_KEYS
            if key in python_series and key in modelica_series and modelica_series[key]
        )
        for key in comparison_keys:
            header.extend(
                [
                    f"python.{key}",
                    f"modelica.{key}",
                    f"abs_error.{key}",
                    f"rel_error.{key}",
                ]
            )
        writer.writerow(header)

        max_abs_errors = {key: 0.0 for key in comparison_keys}
        max_rel_errors = {key: 0.0 for key in comparison_keys}

        for index, time_value in enumerate(python_series["time"]):
            row = [time_value]
            for key in comparison_keys:
                python_value = python_series[key][index]
                modelica_value = _interpolate_series_value(
                    modelica_series["time"],
                    modelica_series[key],
                    time_value,
                )
                abs_error = abs(python_value - modelica_value)
                rel_error = abs_error / max(abs(modelica_value), 1e-9)
                max_abs_errors[key] = max(max_abs_errors[key], abs_error)
                max_rel_errors[key] = max(max_rel_errors[key], rel_error)
                row.extend([python_value, modelica_value, abs_error, rel_error])
            writer.writerow(row)

    summary_lines = []
    for key in comparison_keys:
        summary[key] = (max_abs_errors[key], max_rel_errors[key])
        summary_lines.append(
            f"{key}: max_abs_error={max_abs_errors[key]:.6f}, "
            f"max_rel_error={max_rel_errors[key]:.6%}"
        )
    summary_path.write_text("\n".join(summary_lines) + "\n", encoding="utf-8")
    return comparison_csv_path, summary_path, summary


def export_testmodel_artifacts(
    *,
    output_dir: Path,
    series: dict[str, list[float]],
    modelica_series: dict[str, list[float]] | None = None,
) -> tuple[TestModelArtifacts, dict[str, tuple[float, float]] | None]:
    output_dir.mkdir(parents=True, exist_ok=True)

    primary_csv_path = _write_primary_series_csv(output_dir, series)
    temperature_csv_path = _write_temperature_csv(
        output_dir,
        series["time"],
        series["mytank.T"],
    )
    temperature_svg_path = _write_temperature_svg(
        output_dir,
        series["time"],
        series["mytank.T"],
    )

    comparison_csv_path = None
    comparison_summary_path = None
    comparison_summary = None
    if modelica_series is not None:
        (
            comparison_csv_path,
            comparison_summary_path,
            comparison_summary,
        ) = write_modelica_comparison(output_dir, series, modelica_series)

    return (
        TestModelArtifacts(
            primary_csv_path=primary_csv_path,
            temperature_csv_path=temperature_csv_path,
            temperature_svg_path=temperature_svg_path,
            run_report_path=output_dir / "testmodel_run_report.txt",
            comparison_csv_path=comparison_csv_path,
            comparison_summary_path=comparison_summary_path,
        ),
        comparison_summary,
    )