import os

import numpy as np
import pandas as pd
import scipy.io

import pyocto


class PhaseAssociation:
    """
    Class for phase association using PyOcto.
    """

    def __init__(self, inventory, config):
        """
        Initialize the PhaseAssociation class.

        Args:
            inventory: ObsPy inventory object.
            config (dict): Configuration parameters.
        """
        self._inv = inventory
        self._config = config
        self._setup_velocity_model()
        self._setup_associator()

    def _setup_velocity_model(self):
        """
        Set up the velocity model based on configuration.
        """
        velocity_config = self._config["velocity"]
        if self._config["velocity"]["path"] is None:
            self.velocity_model = pyocto.VelocityModel0D(
                p_velocity=self._config["velocity"]["0D_P_vel"],
                s_velocity=self._config["velocity"]["0D_S_vel"],
                tolerance=self._config["params"]["tolerance"],
                association_cutoff_distance=self._config["params"]["association_cutoff_distance"])
        else:
            if not os.path.exists(velocity_config["path_out"]):
                self._create_1d_vel_model()
            self.velocity_model = pyocto.VelocityModel1D(velocity_config["path_out"],
                                                         tolerance=self._config["params"]["tolerance"])

    def _mat2df(self):
        # Load the .mat file
        mat_contents = scipy.io.loadmat(self._config["velocity"]["path"])
        data = mat_contents["d"][0][0]
        depth = data[0].T[0]
        depth = np.append(depth[::2], depth[-1])
        vel_p = data[1].T[0]
        vel_p = np.append(vel_p[0], vel_p[::2])
        vel_s = data[2].T[0]
        vel_s = np.append(vel_s[0], vel_s[::2])

        # Create DataFrame from the lists
        data = {
            'depth': depth,
            'vp': vel_p,
            'vs': vel_s
        }
        layers = pd.DataFrame(data)

        return layers

    def _create_1d_vel_model(self):

        layers = self._mat2df()
        pyocto.VelocityModel1D.create_model(
            layers,
            self._config["velocity"]["grid_spacing"],
            self._config["velocity"]["grid_hor_extent"],
            self._config["velocity"]["grid_ver_extent"],
            self._config["velocity"]["path_out"]
        )

    def _setup_associator(self):
        """
        Set up the associator based on configuration.
        """
        grid_config = self._config["grid"]
        associator_params = self._config["params"]
        self.associator = pyocto.OctoAssociator.from_area(
            lat=(grid_config["lat_min"], grid_config["lat_max"]),
            lon=(grid_config["lon_min"], grid_config["lon_max"]),
            zlim=(grid_config["z_min"], grid_config["z_max"]),
            time_before=associator_params["time_before"],
            velocity_model=self.velocity_model,
            n_picks=associator_params["n_picks"],
            n_p_picks=associator_params["n_p_picks"],
            n_s_picks=associator_params["n_s_picks"],
            n_p_and_s_picks=associator_params["n_p_and_s_picks"],
            pick_match_tolerance=associator_params["pick_match_tolerance"]
        )

    def associate_phases(self, picks):
        stations = self.associator.inventory_to_df(self._inv)
        events, assignments = self.associator.associate_seisbench(picks, stations)
        events = self.associator.transform_events(events)
        return events, assignments