import os
import logging

import pandas as pd
import scipy.io

from obspy.geodetics import kilometer2degrees
from obspy.taup import TauPyModel
from obspy.taup.taup_create import build_taup_model
from obspy.taup.tau import TauModel


class CustomTauPyModel(TauPyModel):
    """
    Custom TauPyModel for filepath modification.
    """
    def __init__(self, filepath, model="iasp91", verbose=False, planet_flattening=0.0):
        super().__init__(model, verbose, planet_flattening)
        self.verbose = verbose
        self.planet_flattening = planet_flattening
        self.model = self.read_model(filepath)

    @staticmethod
    def read_model(path):
        return TauModel.deserialize(path, cache=None)


class TauPRayTracer:
    """ Class for raytracing seismic waves using TauP """
    def __init__(self, velocity_mat_file):
        logging.basicConfig(filename="application.log",
                        filemode='a',
                        format='%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s',
                        datefmt='%H:%M:%S',
                        level=logging.INFO)
        self.logger = logging.getLogger('converter')
        self.path = "."
        self.velocity_model = self._setup_velocity_model(velocity_mat_file)

    def _setup_velocity_model(self, velocity_mat_file_name):
        velocity_npz_file_path = f"{velocity_mat_file_name[:-3]}npz"
        velocity_npz_file_name = velocity_npz_file_path.split('/')[-1]
        nd_file = self._mat2nd(velocity_mat_file_name)
        build_taup_model(nd_file, output_folder=f"{self.path}")
        os.remove(nd_file)
        try:
            return CustomTauPyModel(filepath=f"{self.path}/{velocity_npz_file_name}")
        except FileNotFoundError as er:
            self.logger.warning(er)

    @staticmethod
    def _mat2nd(velocity_mat_file):
        # Load the .mat file
        mat_contents = scipy.io.loadmat(velocity_mat_file)
        data = mat_contents["d"][0][0]

        # Create DataFrame from the lists
        data_dict = {
            'depth': data[0].T[0],
            'vp': data[1].T[0],
            'vs': data[2].T[0],
            'ro': data[3].T[0],
            'qp': data[4].T[0],
            'qs': data[5].T[0]
        }
        df = pd.DataFrame(data_dict)

        # Adding two new rows to model file - required by TauP
        # Get the last row of the DataFrame
        last_row = df.iloc[-1].copy()
        # Append the last row to the DataFrame

        df = pd.concat([df, pd.DataFrame(last_row).T], ignore_index=True)
        last_row['depth'] = 6178.1
        # Append the modified last row to the DataFrame
        df = pd.concat([df, pd.DataFrame(last_row).T], ignore_index=True)

        # Specify the name of the text file to save
        nd_filename = f'{velocity_mat_file[:-3]}nd'
        # Save the DataFrame to a text file
        df.to_csv(nd_filename, sep='\t', index=False, header=False)

        return nd_filename

    def calculate_arrival(self, distance_in_km, source_depth_in_km, receiver_depth_in_km, phase: str):
        phase_list = [phase, phase.swapcase()]
        return self.velocity_model.get_travel_times(source_depth_in_km=source_depth_in_km,
                                                    distance_in_degree=kilometer2degrees(distance_in_km),
                                                    receiver_depth_in_km=receiver_depth_in_km,
                                                    phase_list=phase_list
                                                    )[0]