from __future__ import annotations  # noqa: F404

from typing import Hashable, List

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

from pypilecore.results.cases_grouper_results import CasesGrouperResults
from pypilecore.results.data_tables import ResultsPandasColumn
from pypilecore.results.result_definitions import CPTResultDefinition
from pypilecore.results.typing import CasesMultiCPTResultsLike


class FigureCPTResultsVersusPtls:
    """
    Interactive figure to show the CPT results of the bearing capacity calculations
    versus the pile tip levels (PTLs).

    The layout of the figure is:
        - X axis: result values.
        - Y axis: pile tip level w.r.t. NAP.
        - Each trace represents a different CPT.

    The figure has a method to switch between cases and results.
    """

    def __init__(self, results_cases: CasesMultiCPTResultsLike) -> None:
        """
        Initializes the figure.

        Parameters
        ----------
        results_cases : CasesMultiCPTResultsProtocol
            The case results of the bearing capacity calculations.
        """
        # Validate the data
        if not isinstance(results_cases, CasesMultiCPTResultsLike):
            raise TypeError(
                f"Incompatible type for 'cases_multi_results': {type(results_cases)}"
            )

        # Set attributes
        self._results = results_cases

        # Initialize the figure
        self._figure = go.FigureWidget()

    @property
    def results(self) -> CasesMultiCPTResultsLike:
        """The results of the bearing capacity calculations."""
        return self._results

    @property
    def data(self) -> pd.DataFrame:
        """The dataframe used to plot the results."""
        return self.results.cpt_results_table.to_pandas()

    @property
    def cases(self) -> List[Hashable]:
        """The case names of each MultiCPTBearingResults."""
        return self.results.cases

    @property
    def test_ids(self) -> List[str]:
        """The test_ids (cpt names) of all the MultiCPTBearingResults."""
        return self.results.test_ids

    @property
    def figure(self) -> go.FigureWidget:
        """The figure widget."""
        return self._figure

    def get_visible_test_ids(self) -> List[go.Scatter]:
        """Returns the visible `test_id` (s) in the figure widget."""
        return [trace.name for trace in self.figure.data if trace.visible is True]

    def show_case_and_result(self, case_name: Hashable, result_name: str) -> None:
        """Shows the results for all CPTs and pile tip levels for the requested `case_name` and `result_name`.

        Parameters
        ----------
        case_name : str
            The name of the case to show.
        result_name : str
            The name of the result to show.

        Raises
        ------
        ValueError
            If the `case_name` is not found in the cases.
            If the `result_name` is not found in the CPTResultDefinitions.
        """
        # Check that case name is in cases.
        if case_name not in self.cases:
            raise ValueError(f"Case name '{case_name}' not found in cases.")

        # Get the result definition that corresponds to the result name.
        result_definition = CPTResultDefinition.get(result_name)

        # Get the visible test_ids
        if len(self.figure.data) == 0:
            # Select all test_ids if there are no traces in the figure yet.
            visible_test_ids = self.test_ids
        else:
            visible_test_ids = self.get_visible_test_ids()

        # Select data for case name and result name.
        mask_case_name = (
            self.data[ResultsPandasColumn.CASE_NAME.value] == case_name
            if case_name is not None
            else self.data[ResultsPandasColumn.CASE_NAME.value].isna()
        )
        selected_data = self.data.loc[
            (mask_case_name)
            & (
                self.data[ResultsPandasColumn.RESULT_NAME.value]
                == result_definition.name
            )
        ]
        traces = []
        for test_id in self.test_ids:
            df = selected_data.loc[
                selected_data[ResultsPandasColumn.TEST_ID.value] == test_id
            ]
            traces.append(
                go.Scatter(
                    x=df[ResultsPandasColumn.RESULT.value],
                    y=df[ResultsPandasColumn.PILE_TIP_LEVEL_NAP.value],
                    mode="lines+markers",
                    name=test_id,
                )
            )

        with self.figure.batch_update():
            # Empty traces
            self.figure.data = []

            # Apply changes
            self.figure.add_traces(traces)

            calculation_type = (
                "Grouper" if isinstance(self.results, CasesGrouperResults) else ""
            )

            self.figure.update_layout(
                title="CPT "
                + calculation_type
                + f"Results vs. Pile tip level<br>Case: {case_name}, Result: {result_definition.value.html}",
                xaxis=go.layout.XAxis(
                    title=f"{result_definition.value.html} [{result_definition.value.unit}]",
                    title_font_size=18,
                ),
                showlegend=True,
                # The following parameters shouldn't be necessary to update
                # (and therefore could be assigned on initialization of self.figure),
                # but they are needed to avoid a bug in plotly since dash > 3.0.
                height=800,
                width=800,
                legend_title="CPT",
                colorway=px.colors.qualitative.Plotly,
                yaxis=go.layout.YAxis(
                    title="Pile tip level [m NAP]",
                    title_font_size=18,
                ),
                autosize=False,
            )

            self.figure.update_traces(
                selector=lambda x: x.name in visible_test_ids,
                patch=dict(
                    visible=True,
                ),
            )

            self.figure.update_traces(
                selector=lambda x: x.name not in visible_test_ids,
                patch=dict(
                    visible="legendonly",
                ),
            )
