import pandas as pd

from obspy.core.event import Catalog
from obspy.core.event import Event
from obspy.core.event import Pick
from obspy.core.event import Origin
from obspy.core.event import OriginQuality

from obspy import UTCDateTime
from obspy.core.event.base import WaveformStreamID, Comment


class CatalogConverter:
    """
    Class for converting SeisBench and pyOcto detection results to ObsPy Catalog, which can be saved as QuakeML.
    """
    def __init__(self, config, picks, catalog_df, assignments_df, name_of_algorithm):
        self._catalog_df = catalog_df
        self._assignments_df = assignments_df
        self._picks = picks
        self._config = config
        self._name_of_algorithm = name_of_algorithm
        self.catalog = None

    def _convert_pick(self, p):
        """
        Function converts picks from SeisBench to ObsPy format
        :param p: SeisBench pick
        :return: ObsPy pick
        """
        pick = Pick()
        pick.time = UTCDateTime(p.peak_time.datetime)
        pick.waveform_id = WaveformStreamID(network_code=p.trace_id.split(".")[0],
                                            station_code=p.trace_id.split(".")[1],
                                            channel_code=p.trace_id.split(".")[2])
        if p.phase == 'P':
            pick.phase_hint = self._config["P_hint"]
        elif p.phase == 'S':
            pick.phase_hint = self._config["S_hint"]
        pick.evaluation_mode = 'automatic'
        pick.evaluation_status = 'preliminary'
        return pick

    def _convert_origin(self, origin_sb, list_of_picks_sb):
        origin = Origin()
        origin.time = UTCDateTime(pd.to_datetime(origin_sb.time, unit='s').to_pydatetime())
        origin.latitude = origin_sb.latitude  # float
        origin.longitude = origin_sb.longitude  # float
        origin.depth = origin_sb.depth  # float in kilometers (SWIP5 origin version) down the see level
        origin.depth_type = 'operator assigned'
        # TODO: make sure that region is not necessary
        # origin.region = self._config["region"]
        origin.evaluation_mode = "automatic"
        origin.evaluation_status = 'preliminary'
        origin.comments.append(Comment(text=f"Localized by: {self._name_of_algorithm}", force_resource_id=False))
        origin.quality = OriginQuality(used_phase_count=len(list_of_picks_sb))
        return origin

    def _convert_event(self, origin_sb, list_of_picks_sb):
        """
        Function convert GaMMa detection to ObsPy Event
        :param origin_sb:
        :param list_of_picks_sb:
        :return:
        """
        event = Event()
        for p in list_of_picks_sb:
            pick = self._convert_pick(p)
            event.picks.append(pick)
        origin = self._convert_origin(origin_sb, list_of_picks_sb)
        event.origins.append(origin)
        return event

    @staticmethod
    def _append_pick_trace_id(pick, stream):
        """
        Function assigns channel to pick - it is useful for work with SWIP
        :param pick:
        :param stream:
        :return:
        """
        channel = stream[0].stats.channel
        if pick.phase == "P":
            pick.trace_id = pick.trace_id + channel[:-1] + "Z"
        if pick.phase == "S":
            pick.trace_id = pick.trace_id + channel[:-1] + "E"
        return pick

    def catalog2obspy(self):
        """
        Function convert GaMMa catalog and SeisBench picks
        :return: ObsPy Catalog object
        """
        # TODO: make sure that resource id is necessary
        #cat = Catalog(resource_id=self._config["resource_id"])
        cat = Catalog()
        for j, row in self._catalog_df.iterrows():
            event = self._catalog_df.iloc[j]
            event_picks = [self._picks[i] for i in
                           self._assignments_df[self._assignments_df["event_idx"] ==
                                                event["idx"]]["pick_idx"]]
            event_obspy = self._convert_event(event, event_picks)
            cat.append(event_obspy)
        self.catalog = cat

    def save_catalog_to_file(self, file_path):
        """
        Save ObsPy catalog to a file.

        Args:
            file_path (str): The file path where the catalog will be saved.

        Returns:
            None
        """
        try:
            self.catalog.write(file_path, format="QUAKEML")
            print(f"Catalog saved successfully to {file_path}")
        except Exception as e:
            print(f"Error occurred while saving catalog: {e}")