from datetime import datetime
from pathlib import Path
import xarray as xr 

from matplotlib import pyplot as plt
from matplotlib.figure import Figure
import numpy as np

from gfatpy.radar.rpg_nc import rpg
from gfatpy.utils.utils import parse_datetime


class nebula():
    def __init__(self, ka_nc_path: Path, ww_nc_path: Path):
        self.ka = rpg(ka_nc_path)
        self.ww = rpg(ww_nc_path)        
        self._level = None
        self.nb = self.ka
        self.nb._data = self.combined_data()

    @property
    def level(self) -> int:
        if self.ka.level != self.ww.level:
            raise ValueError("Ka- and W-band measurement levels do not match.")        
        self._level = self.ka.level
        return self._level

    def combined_data(self) -> xr.Dataset:
        if self.level == 1:
            if self.nb._data is None:
                self.nb._data["DWR"] = 10 * np.log10(self.ka.data["Ze"] / self.ww.data["Ze"])
                self.nb._data["DWR"].attrs["units"] = "dB"
                self.nb._data["DWR"].attrs["long_name"] = "Ka-W DWR"

                self.nb._data["DDV"] = self.ka.data["v"] - self.ww.data["v"]
                self.nb._data["DDV"].attrs["units"] = "m/s"
                self.nb._data["DDV"].attrs["long_name"] = "Ka-W DDV"
        return self.nb._data

    def quicklook(self, variable: str | list[str] | None = None, **kwargs):
        self.ka.quicklook(variable, **kwargs)
        self.ww.quicklook(variable, **kwargs)

    def plot_profile(
        self,
        target_time: datetime | np.datetime64,
        range_limits: float,
        variable: str | list[str] | None = None,
        **kwargs,
    ) -> tuple[Figure, Path | None]:
        _kwargs = kwargs.copy()        
        _kwargs["savefig"] = False
        _kwargs["color_list"] = [kwargs.get("ka_color", "red")]

        fig1, _ = self.ka.plot_profile(
            target_times=target_time,
            range_limits=range_limits,
            variable=variable,
            **_kwargs,
        )

        _kwargs["fig"] = fig1
        _kwargs["color_list"] = [kwargs.get("w_color", "blue")]
        fig, _ = self.ww.plot_profile(
            target_times=target_time,
            range_limits=range_limits,
            variable=variable,
            **_kwargs,
        )

        ka_line, w_line = fig.findobj(plt.Line2D)[0], fig.findobj(plt.Line2D)[1]
        ka_line.set_label(f"Ka-band")
        w_line.set_label(f"W-band")

        ax = fig.get_axes()[0]
        ax.legend()

        if kwargs.get("savefig", False):
            filepath = (
                kwargs.get("output_dir", Path.cwd())
                / f"profile_{variable}_{target_time:%Y%m%dT%H%M%S}.png"
            )
            fig.savefig(filepath, dpi=kwargs.get("dpi", 300))
        return fig, filepath

    def plot_timeseries(
        self,
        target_range: float,
        time_limits: tuple[datetime | np.datetime64, datetime | np.datetime64] | None = None,
        variable: str | list[str] | None = None,
        **kwargs,
    ) -> tuple[Figure, Path | None]:
        
        _kwargs = kwargs.copy()        
        _kwargs["savefig"] = False
        _kwargs["color_list"] = [kwargs.get("ka_color", "red")]

        if  time_limits is None:
            time_limits = (parse_datetime(self.ka.data.time[0].values), parse_datetime(self.ka.data.time[-1].values))

        fig1, _ = self.ka.plot_timeseries(
            target_ranges=target_range,
            time_limits=time_limits,
            variable=variable,
            **_kwargs,
        )

        _kwargs["fig"] = fig1
        _kwargs["color_list"] = [kwargs.get("w_color", "blue")]
        fig, _ = self.ww.plot_timeseries(
            target_ranges=target_range,
            time_limits=time_limits,
            variable=variable,
            **_kwargs,
        )

        ka_line, w_line = fig.findobj(plt.Line2D)[0], fig.findobj(plt.Line2D)[1]
        ka_line.set_label(f"Ka-band")
        w_line.set_label(f"W-band")

        ax = fig.get_axes()[0]
        variable_string = self.ka.data[variable].attrs['long_name'].replace('W-band','').replace("Ka-band",'')
        ax.set_ylabel(f"{variable_string}, [{self.ka.data[variable].attrs['units']}]")
        ax.legend()

        if kwargs.get("savefig", False):
            filepath = (
                kwargs.get("output_dir", Path.cwd())
                / f"timeseries_{variable}_{time_limits[0]:%Y%m%dT%H%M%S}_{time_limits[-1]:%Y%m%dT%H%M%S}_{target_range}.png"
            )
            fig.savefig(filepath, dpi=kwargs.get("dpi", 300))
        return fig, filepath


    def plot_spectrum(
        self, target_time: datetime | np.datetime64, target_range: float, **kwargs
    ) -> tuple[Figure, Path | None]:
        _kwargs = kwargs.copy()
        _kwargs.pop("output_dir")
        _kwargs.pop("savefig")
        _kwargs["color"] = kwargs.get("ka_color", "red")
        fig, _ = self.ka.plot_spectrum(target_time, target_range, **_kwargs)

        _kwargs["fig"] = fig
        _kwargs["color"] = kwargs.get("w_color", "blue")
        fig, _ = self.ww.plot_spectrum(target_time, target_range, **_kwargs)
        ka_line, w_line = fig.findobj(plt.Line2D)[0], fig.findobj(plt.Line2D)[2]
        ka_line.set_label(f"Ka-band {ka_line.get_label()}")
        w_line.set_label(f"W-band {w_line.get_label()}")

        ax = fig.get_axes()[0]
        ax.legend()

        if kwargs.get("savefig", False):
            filepath = (
                kwargs.get("output_dir", Path.cwd())
                / f"{target_time:%Y%m%dT%H%M%S}_{target_range:.0f}_spectrum.png"
            )
            fig.savefig(filepath, dpi=kwargs.get("dpi", 300))

        return fig, filepath
