from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from app.simulation.examples.test_mql.primitives.pneumatic import m3_to_cm3
from app.simulation.reporting.amesim_results import AmesimResults, load_test_mql_amesim_results
from app.simulation.reporting.test_mql_comparison import TestMqlComparisonResult
from app.simulation.reporting.test_mql_observations import (
    TestMqlObservationCatalog,
    build_test_mql_observation_catalog,
)
from app.simulation.reporting.test_mql_output_schema import (
    TestMqlOutputSchema,
    build_test_mql_output_schema,
)
from app.simulation.reporting.test_mql_output_validation import (
    TestMqlValidatedOutput,
    compare_validated_test_mql_output,
    validate_test_mql_output,
)
from app.simulation.reporting.test_mql_variables import build_test_mql_variable_catalog
from app.simulation.examples.test_mql.mechanical import (
    TestMqlMechanicalAssembly,
    build_test_mql_mechanical_assembly,
)
from app.simulation.examples.test_mql.pneumatic import (
    TestMqlPneumaticAssembly,
    build_test_mql_pneumatic_assembly,
)


@dataclass(frozen=True)
class TestMqlComputedPistonGeometryRun:
    amesim_results: AmesimResults
    observation_catalog: TestMqlObservationCatalog
    output_schema: TestMqlOutputSchema
    mechanical_assembly: TestMqlMechanicalAssembly
    output: TestMqlValidatedOutput
    comparison: TestMqlComparisonResult

    @property
    def sample_count(self) -> int:
        return len(self.output.times)

    @property
    def signal_count(self) -> int:
        return len(self.output.data_paths)


@dataclass(frozen=True)
class TestMqlComputedGeometryRun:
    amesim_results: AmesimResults
    observation_catalog: TestMqlObservationCatalog
    output_schema: TestMqlOutputSchema
    mechanical_assembly: TestMqlMechanicalAssembly
    pneumatic_assembly: TestMqlPneumaticAssembly
    output: TestMqlValidatedOutput
    comparison: TestMqlComparisonResult

    @property
    def sample_count(self) -> int:
        return len(self.output.times)

    @property
    def signal_count(self) -> int:
        return len(self.output.data_paths)


@dataclass(frozen=True)
class TestMqlComputedLineRelationsRun:
    amesim_results: AmesimResults
    observation_catalog: TestMqlObservationCatalog
    output_schema: TestMqlOutputSchema
    output: TestMqlValidatedOutput
    comparison: TestMqlComparisonResult

    @property
    def sample_count(self) -> int:
        return len(self.output.times)

    @property
    def signal_count(self) -> int:
        return len(self.output.data_paths)


@dataclass(frozen=True)
class TestMqlComputedPneumaticRelationsRun:
    amesim_results: AmesimResults
    observation_catalog: TestMqlObservationCatalog
    output_schema: TestMqlOutputSchema
    output: TestMqlValidatedOutput
    comparison: TestMqlComparisonResult

    @property
    def sample_count(self) -> int:
        return len(self.output.times)

    @property
    def signal_count(self) -> int:
        return len(self.output.data_paths)


@dataclass(frozen=True)
class TestMqlComputedMechanicalRelationsRun:
    amesim_results: AmesimResults
    observation_catalog: TestMqlObservationCatalog
    output_schema: TestMqlOutputSchema
    mechanical_assembly: TestMqlMechanicalAssembly
    output: TestMqlValidatedOutput
    comparison: TestMqlComparisonResult

    @property
    def sample_count(self) -> int:
        return len(self.output.times)

    @property
    def signal_count(self) -> int:
        return len(self.output.data_paths)


def run_test_mql_computed_piston_geometry(
    archive_path: Path,
) -> TestMqlComputedPistonGeometryRun:
    amesim_results = load_test_mql_amesim_results(archive_path)
    observation_catalog = build_test_mql_observation_catalog(amesim_results)
    output_schema = build_test_mql_output_schema(
        amesim_results,
        observation_catalog=observation_catalog,
    )
    variable_catalog = build_test_mql_variable_catalog(amesim_results)
    mechanical_assembly = build_test_mql_mechanical_assembly(
        amesim_results=amesim_results,
        variable_catalog=variable_catalog,
    )
    output_series = _compute_piston_geometry_series(amesim_results, mechanical_assembly)
    output_data_paths = tuple(output_series)
    output = validate_test_mql_output(
        times=amesim_results.times,
        series_by_data_path=output_series,
        schema=output_schema,
        data_paths=output_data_paths,
    )
    comparison = compare_validated_test_mql_output(
        times=output.times,
        series_by_data_path=output.series_by_data_path,
        schema=output_schema,
        amesim_results=amesim_results,
        data_paths=output.data_paths,
    )
    return TestMqlComputedPistonGeometryRun(
        amesim_results=amesim_results,
        observation_catalog=observation_catalog,
        output_schema=output_schema,
        mechanical_assembly=mechanical_assembly,
        output=output,
        comparison=comparison,
    )


def run_test_mql_computed_geometry(
    archive_path: Path,
) -> TestMqlComputedGeometryRun:
    amesim_results = load_test_mql_amesim_results(archive_path)
    observation_catalog = build_test_mql_observation_catalog(amesim_results)
    output_schema = build_test_mql_output_schema(
        amesim_results,
        observation_catalog=observation_catalog,
    )
    variable_catalog = build_test_mql_variable_catalog(amesim_results)
    mechanical_assembly = build_test_mql_mechanical_assembly(
        amesim_results=amesim_results,
        variable_catalog=variable_catalog,
    )
    pneumatic_assembly = build_test_mql_pneumatic_assembly()
    output_series = {
        **_compute_piston_geometry_series(amesim_results, mechanical_assembly),
        **_compute_variable_chamber_volume_series(
            amesim_results,
            mechanical_assembly,
            pneumatic_assembly,
        ),
    }
    output_data_paths = tuple(output_series)
    output = validate_test_mql_output(
        times=amesim_results.times,
        series_by_data_path=output_series,
        schema=output_schema,
        data_paths=output_data_paths,
    )
    comparison = compare_validated_test_mql_output(
        times=output.times,
        series_by_data_path=output.series_by_data_path,
        schema=output_schema,
        amesim_results=amesim_results,
        data_paths=output.data_paths,
    )
    return TestMqlComputedGeometryRun(
        amesim_results=amesim_results,
        observation_catalog=observation_catalog,
        output_schema=output_schema,
        mechanical_assembly=mechanical_assembly,
        pneumatic_assembly=pneumatic_assembly,
        output=output,
        comparison=comparison,
    )


def run_test_mql_computed_line_relations(
    archive_path: Path,
) -> TestMqlComputedLineRelationsRun:
    amesim_results = load_test_mql_amesim_results(archive_path)
    observation_catalog = build_test_mql_observation_catalog(amesim_results)
    output_schema = build_test_mql_output_schema(
        amesim_results,
        observation_catalog=observation_catalog,
    )
    output_series = _compute_line_reversed_series(amesim_results, observation_catalog)
    output_data_paths = tuple(output_series)
    output = validate_test_mql_output(
        times=amesim_results.times,
        series_by_data_path=output_series,
        schema=output_schema,
        data_paths=output_data_paths,
    )
    comparison = compare_validated_test_mql_output(
        times=output.times,
        series_by_data_path=output.series_by_data_path,
        schema=output_schema,
        amesim_results=amesim_results,
        data_paths=output.data_paths,
    )
    return TestMqlComputedLineRelationsRun(
        amesim_results=amesim_results,
        observation_catalog=observation_catalog,
        output_schema=output_schema,
        output=output,
        comparison=comparison,
    )


def run_test_mql_computed_pneumatic_relations(
    archive_path: Path,
) -> TestMqlComputedPneumaticRelationsRun:
    amesim_results = load_test_mql_amesim_results(archive_path)
    observation_catalog = build_test_mql_observation_catalog(amesim_results)
    output_schema = build_test_mql_output_schema(
        amesim_results,
        observation_catalog=observation_catalog,
    )
    output_series = {
        **_compute_chamber_duplicate_series(amesim_results, observation_catalog),
        **_compute_orifice_reversed_series(amesim_results, observation_catalog),
    }
    output_data_paths = tuple(output_series)
    output = validate_test_mql_output(
        times=amesim_results.times,
        series_by_data_path=output_series,
        schema=output_schema,
        data_paths=output_data_paths,
    )
    comparison = compare_validated_test_mql_output(
        times=output.times,
        series_by_data_path=output.series_by_data_path,
        schema=output_schema,
        amesim_results=amesim_results,
        data_paths=output.data_paths,
    )
    return TestMqlComputedPneumaticRelationsRun(
        amesim_results=amesim_results,
        observation_catalog=observation_catalog,
        output_schema=output_schema,
        output=output,
        comparison=comparison,
    )


def run_test_mql_computed_mechanical_relations(
    archive_path: Path,
) -> TestMqlComputedMechanicalRelationsRun:
    amesim_results = load_test_mql_amesim_results(archive_path)
    observation_catalog = build_test_mql_observation_catalog(amesim_results)
    output_schema = build_test_mql_output_schema(
        amesim_results,
        observation_catalog=observation_catalog,
    )
    variable_catalog = build_test_mql_variable_catalog(amesim_results)
    mechanical_assembly = build_test_mql_mechanical_assembly(
        amesim_results=amesim_results,
        variable_catalog=variable_catalog,
    )
    output_series = {
        **_compute_mass_duplicate_series(amesim_results, mechanical_assembly),
        **_compute_inactive_mass_force_series(amesim_results, mechanical_assembly),
        **_compute_zero_force_source_series(amesim_results, mechanical_assembly),
    }
    output_data_paths = tuple(output_series)
    output = validate_test_mql_output(
        times=amesim_results.times,
        series_by_data_path=output_series,
        schema=output_schema,
        data_paths=output_data_paths,
    )
    comparison = compare_validated_test_mql_output(
        times=output.times,
        series_by_data_path=output.series_by_data_path,
        schema=output_schema,
        amesim_results=amesim_results,
        data_paths=output.data_paths,
    )
    return TestMqlComputedMechanicalRelationsRun(
        amesim_results=amesim_results,
        observation_catalog=observation_catalog,
        output_schema=output_schema,
        mechanical_assembly=mechanical_assembly,
        output=output,
        comparison=comparison,
    )


def _compute_piston_geometry_series(
    amesim_results: AmesimResults,
    mechanical_assembly: TestMqlMechanicalAssembly,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for alias in sorted(mechanical_assembly.pistons):
        piston = mechanical_assembly.pistons[alias]
        geometry = piston.geometry()
        x4 = amesim_results.series(f"x4@{alias}")
        x5 = amesim_results.series(f"x5@{alias}")
        v4 = amesim_results.series(f"v4@{alias}")
        v5 = amesim_results.series(f"v5@{alias}")
        series_by_data_path[f"length@{alias}"] = tuple(
            geometry.chamber_length_mm(port4, port5)
            for port4, port5 in zip(x4, x5)
        )
        series_by_data_path[f"vol1@{alias}"] = tuple(
            geometry.chamber_volume_cm3(port4, port5)
            for port4, port5 in zip(x4, x5)
        )
        series_by_data_path[f"vvol1@{alias}"] = tuple(
            geometry.chamber_volume_rate_l_min(port4, port5)
            for port4, port5 in zip(v4, v5)
        )
    return series_by_data_path


def _compute_variable_chamber_volume_series(
    amesim_results: AmesimResults,
    mechanical_assembly: TestMqlMechanicalAssembly,
    pneumatic_assembly: TestMqlPneumaticAssembly,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for chamber_alias in sorted(pneumatic_assembly.variable_chambers):
        chamber = pneumatic_assembly.variable_chambers[chamber_alias]
        piston_alias = _piston_alias_for_variable_chamber(chamber_alias)
        piston = mechanical_assembly.pistons[piston_alias]
        geometry = piston.geometry()
        x4 = amesim_results.series(f"x4@{piston_alias}")
        x5 = amesim_results.series(f"x5@{piston_alias}")
        dead_volume_cm3 = m3_to_cm3(chamber.dead_volume)
        series_by_data_path[f"vol@{chamber_alias}"] = tuple(
            dead_volume_cm3 + geometry.chamber_volume_cm3(port4, port5)
            for port4, port5 in zip(x4, x5)
        )
    return series_by_data_path


def _piston_alias_for_variable_chamber(chamber_alias: str) -> str:
    if not chamber_alias.startswith("pn_c1"):
        raise ValueError(f"Unexpected PNCH012 alias: {chamber_alias}")
    return chamber_alias.replace("pn_c1", "pn_brp2", 1)


def _compute_mass_duplicate_series(
    amesim_results: AmesimResults,
    mechanical_assembly: TestMqlMechanicalAssembly,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for alias in sorted(mechanical_assembly.masses):
        for signal_name in ("x1", "v1", "acc1"):
            source_path = f"{signal_name}@{alias}"
            duplicate_path = f"{signal_name}dup@{alias}"
            series_by_data_path[duplicate_path] = tuple(
                -value for value in amesim_results.series(source_path)
            )
    return series_by_data_path


def _compute_inactive_mass_force_series(
    amesim_results: AmesimResults,
    mechanical_assembly: TestMqlMechanicalAssembly,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for alias in sorted(mechanical_assembly.masses):
        mass = mechanical_assembly.masses[alias].endstop()
        x1 = amesim_results.series(f"x1@{alias}")
        v1 = amesim_results.series(f"v1@{alias}")
        series_by_data_path[f"Fmin@{alias}"] = tuple(
            mass.lower_static_force_magnitude(displacement)
            for displacement in x1
        )
        series_by_data_path[f"Fvisc@{alias}"] = tuple(
            mass.viscous_friction_force(velocity)
            for velocity in v1
        )
        series_by_data_path[f"Ffric@{alias}"] = tuple(0.0 for _ in x1)
    return series_by_data_path


def _compute_zero_force_source_series(
    amesim_results: AmesimResults,
    mechanical_assembly: TestMqlMechanicalAssembly,
) -> dict[str, tuple[float, ...]]:
    return {
        f"fzero@{alias}": tuple(0.0 for _ in amesim_results.times)
        for alias in sorted(mechanical_assembly.zero_force_sources)
    }


def _compute_chamber_duplicate_series(
    amesim_results: AmesimResults,
    observation_catalog: TestMqlObservationCatalog,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for binding in observation_catalog.chambers.bindings:
        pressure_series = amesim_results.series(binding.pressure_path)
        temperature_series = amesim_results.series(binding.temperature_path)
        for duplicate_path in binding.pressure_duplicate_paths:
            series_by_data_path[duplicate_path] = tuple(pressure_series)
        for duplicate_path in binding.temperature_duplicate_paths:
            series_by_data_path[duplicate_path] = tuple(temperature_series)
    return series_by_data_path


def _compute_orifice_reversed_series(
    amesim_results: AmesimResults,
    observation_catalog: TestMqlObservationCatalog,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for binding in observation_catalog.orifices.bindings:
        series_by_data_path[binding.reversed_mass_flow_path] = tuple(
            -value for value in amesim_results.series(binding.primary_mass_flow_path)
        )
        series_by_data_path[binding.reversed_enthalpy_flow_path] = tuple(
            -value for value in amesim_results.series(binding.primary_enthalpy_flow_path)
        )
    return series_by_data_path


def _compute_line_reversed_series(
    amesim_results: AmesimResults,
    observation_catalog: TestMqlObservationCatalog,
) -> dict[str, tuple[float, ...]]:
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for binding in observation_catalog.lines.by_submodel("PNL00R"):
        if len(binding.mass_flow_paths) != 2 or len(binding.enthalpy_flow_paths) != 2:
            raise ValueError(f"Expected two PNL00R flow paths for {binding.alias}.")
        primary_mass_path, reversed_mass_path = binding.mass_flow_paths
        primary_enthalpy_path, reversed_enthalpy_path = binding.enthalpy_flow_paths
        series_by_data_path[reversed_mass_path] = tuple(
            -value for value in amesim_results.series(primary_mass_path)
        )
        series_by_data_path[reversed_enthalpy_path] = tuple(
            -value for value in amesim_results.series(primary_enthalpy_path)
        )
    return series_by_data_path