from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from PythonModels.components.amesim_pneumatic import m3_to_cm3
from PythonModels.reporting.amesim_results import AmesimResults, load_test_mql_amesim_results
from PythonModels.reporting.test_mql_comparison import TestMqlComparisonResult
from PythonModels.reporting.test_mql_observations import (
    TestMqlObservationCatalog,
    build_test_mql_observation_catalog,
)
from PythonModels.reporting.test_mql_output_schema import (
    TestMqlOutputSchema,
    build_test_mql_output_schema,
)
from PythonModels.reporting.test_mql_output_validation import (
    TestMqlValidatedOutput,
    compare_validated_test_mql_output,
    validate_test_mql_output,
)
from PythonModels.reporting.test_mql_variables import build_test_mql_variable_catalog
from PythonModels.systems.test_mql_mechanical import (
    TestMqlMechanicalAssembly,
    build_test_mql_mechanical_assembly,
)
from PythonModels.systems.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 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_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)
    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