# -*- coding: utf-8 -*-

import os

import numpy as np
import tilupy.read

# Classify results
AVAILABLE_OUTPUT = ["h", "hvert", "u", "hu", "hu2", "ux", "uy", "hux", "huy"]
"""All output available for a shaltop simulation.

Implemented states:

    - h : Flow thickness (normal to the surface)
    - hux : Momentum flux component in X direction
    - huy : Momentum flux component in Y direction
    
Output computed from other output:
    
    - hvert : True vertical flow thickness
    - ux : Velocity component in X direction
    - uy : Velocity component in Y direction
    - u : Norm of the velocity (from ux and uy)
    - hu : Momentum flux (from h and u)
    - hu2 : Convective momentum flux (from h and u)
"""


def extract_saval_ascii(path: str) -> np.ndarray:
    """Extract ascii file data generated by saval2D.

    Parameters
    ----------
    path : str
        Path to the ascii file.

    Returns
    -------
    numpy.ndarray
        Values in the ascii file.
    """
    with open(path, "r") as f:
        lines = f.readlines()
    start_index = 6

    h_str = lines[start_index:]

    h = [list(map(float, line.split())) for line in h_str if line.strip()]
    return np.array(h)


class Results(tilupy.read.Results):
    """Results of saval2D simulations.

    This class is the results class for saval2D. Reading results from saval2D outputs
    are done in this class.

    This class has all the global and quick attributes of the parent class. The quick
    attributes are only computed if needed and can be deleted to clean memory.

    In addition to these attributes, there are those necessary for the operation of
    reading the saval2D results.

    Parameters
    ----------
    folder : str
        Path to the folder containing the simulation files.
    raster_topo : str
        Topographic raster name.

    Attributes
    ----------
        _simu_inputs : str
            Path to the input folder of saval2D simulation.
        _simu_outputs : str
            Path to the output folder of saval2D simulation.
        _raster : str
            Raster file name.
        _params : dict
            Dictionary storing all simulation parameters.
    """

    def __init__(self, folder, raster_topo="mntsimulation.asc"):
        super().__init__()
        self._code = "saval2D"

        self._folder = folder
        self._folder_output = folder

        self._simu_inputs = (
            os.path.join(folder, "inputs")
            if os.path.exists(os.path.join(folder, "inputs"))
            else None
        )
        self._simu_outputs = (
            os.path.join(folder, "outputs")
            if os.path.exists(os.path.join(folder, "outputs"))
            else None
        )

        if not raster_topo.endswith(".asc"):
            raster_topo = raster_topo + ".asc"
        self._raster = raster_topo

        self._params = dict()

        if self._simu_inputs is not None:
            self._x, self._y, self._zinit = tilupy.raster.read_ascii(
                os.path.join(self._simu_inputs, self._raster)
            )
        self._nx, self._ny = len(self._x), len(self._y)
        self._dx = self._x[1] - self._x[0]
        self._dy = self._y[1] - self._y[0]

        self._params["nx"] = len(self.x)
        self._params["ny"] = len(self.y)

        # Rheology
        rheol = self.get_rheological_parameters(
            os.path.join(self._simu_outputs, "log.txt")
        )
        self._params["tau/rho"] = rheol[0]
        self._params["mu"] = rheol[1]
        self._params["xi"] = rheol[2]

        # Create self._tim
        self._extract_output("X")

    def get_rheological_parameters(self, path_to_log: str) -> list[str]:
        """Get rheological parameters of a saval2D simulation.

        Parameters
        ----------
        path_to_log : str
            Path to the log file containing the value of the rheological parameters.

        Returns
        -------
        list[str]
            List of rheological parameters : tau/rho, mu and xi.
        """
        values = None
        with open(path_to_log, "r") as f:
            for line in f:
                line = line.strip()
                try:
                    parts = list(map(float, line.split()))
                    if len(parts) == 3:
                        values = parts
                        break
                except ValueError:
                    continue
        return values

    def get_times(self, path: str) -> list[float]:
        """Get simulation time steps.

        Parameters
        ----------
        path : str
            Path to the log file.

        Returns
        -------
        list[float]
            List of recorded time steps.
        """
        t_list = []
        with open(path, "r") as f:
            lines = f.readlines()

        start_index = 16
        for line in lines[start_index:]:
            if line.strip() == "":
                continue
            else:
                _, t, _, _, _ = line.split()
                t_list.append(float(t))

        return t_list

    def _extract_output(
        self, name: str, **kwargs
    ) -> tilupy.read.TemporalResults2D | tilupy.read.AbstractResults:
        """Result extraction for saval2D files.

        Parameters
        ----------
        name : str
            Wanted output. Can access to variables in :data:`AVAILABLE_OUTPUT`.

        Returns
        -------
        tilupy.read.TemporalResults2D | tilupy.read.AbstractResults
            Wanted output. If no output computed, return an object of :class:`tilupy.read.AbstractResults`.
        """
        # Read thicknesses or velocity components
        d = None
        t = None
        notation = None

        tim = [0]
        t_list = self.get_times(os.path.join(self._simu_outputs, "log.txt"))

        _, _, h_init = tilupy.raster.read_ascii(
            os.path.join(self._simu_inputs, "zdepsimulation.asc")
        )

        h_list = [h_init]

        ux_list = []
        uy_list = []
        u_list = []
        ux_list.append(np.zeros_like(h_init))
        uy_list.append(np.zeros_like(h_init))
        u_list.append(np.zeros_like(h_init))

        qx_list = []
        qy_list = []
        q_list = []
        qx_list.append(np.zeros_like(h_init))
        qy_list.append(np.zeros_like(h_init))
        q_list.append(np.zeros_like(h_init))

        for T in range(len(t_list)):
            h_t = extract_saval_ascii(
                os.path.join(self._simu_outputs, f"resuh{T + 1}.asc")
            )
            qu_t = extract_saval_ascii(
                os.path.join(self._simu_outputs, f"resuqu{T + 1}.asc")
            )
            qv_t = extract_saval_ascii(
                os.path.join(self._simu_outputs, f"resuqv{T + 1}.asc")
            )

            # Condition to avoid flow velocity on part with no fluid
            h_t[h_t < 0.0001] = 0

            ux_t = np.divide(qu_t, h_t, out=np.zeros_like(qu_t), where=h_t != 0)
            uy_t = np.divide(qv_t, h_t, out=np.zeros_like(qv_t), where=h_t != 0)

            u_t = np.sqrt(ux_t**2 + uy_t**2)
            q_t = np.sqrt(qu_t**2 + qv_t**2)

            h_list.append(h_t)

            ux_list.append(ux_t)
            uy_list.append(uy_t)
            u_list.append(u_t)

            qx_list.append(qu_t)
            qy_list.append(qv_t)
            q_list.append(q_t)

            tim.append(t_list[T])

        if self._tim is None:
            self._tim = np.array(tim)

        extracted_outputs = {
            "h": np.stack(h_list, axis=-1),
            "u": np.stack(u_list, axis=-1),
            "hu": np.stack(q_list, axis=-1),
            "ux": np.stack(ux_list, axis=-1),
            "uy": np.stack(uy_list, axis=-1),
            "hux": np.stack(qx_list, axis=-1),
            "huy": np.stack(qy_list, axis=-1),
        }

        if name in ["h", "u", "hu", "ux", "uy", "hux", "huy"]:
            d = extracted_outputs[name]
            t = self._tim

        if name == "hu2":
            d = extracted_outputs["hu"] * extracted_outputs["u"]
            t = self._tim

        elif name == "hvert":
            if self._costh is None:
                self._costh = self.compute_costh()
            d = extracted_outputs["h"] / self._costh[:, :, np.newaxis]
            t = self._tim

        # if name == "ek":
        #     if self._costh is None:
        #             self._costh = self.compute_costh()

        #     d = []
        #     for i in range(len(self._tim)):
        #         d.append(np.sum((np.stack(available_outputs["hu"], axis=-1)[:, :, i] *
        #                          np.stack(available_outputs["u"], axis=-1)[:, :, i] *
        #                          self._dx *
        #                          self._dy)
        #                         / self._costh[:, :]))
        #     d = np.array(d)
        #     t = self._tim

        if t is None:
            return tilupy.read.AbstractResults(name, d, notation=notation)

        else:
            if d.ndim == 3:
                return tilupy.read.TemporalResults2D(
                    name, d, t, notation=notation, x=self._x, y=self._y, z=self._zinit
                )
        return None

    def _read_from_file(self, *args, **kwargs):
        """Not useful"""
        return "No _read_from_file for Saval2D."
