from __future__ import annotations

import hashlib
import json
import logging
from pathlib import Path

import numpy as np
import xarray as xr

from sasktran2.optical.database import (
    OpticalDatabaseGenericAbsorber,
)
from sasktran2.optical.hitran import HITRANAbsorber
from sasktran2.units import wavenumber_cminv_to_wavlength_nm
from sasktran2.util import get_hapi

from .base import CachedDatabase
from .hitran_line import HITRANLineDatabase

PRESSURE_GRID = (
    np.array(
        [
            1.03181655e03,
            9.99910226e02,
            9.72231997e02,
            9.44960640e02,
            9.18098202e02,
            8.91646636e02,
            8.65607893e02,
            8.39983721e02,
            8.14775620e02,
            7.89985094e02,
            7.65613488e02,
            7.41661804e02,
            7.18131094e02,
            6.95022257e02,
            6.72335891e02,
            6.50072549e02,
            6.28232578e02,
            6.06816233e02,
            5.85823558e02,
            5.65254408e02,
            5.45108580e02,
            5.25385676e02,
            5.06085044e02,
            4.87205987e02,
            4.68747652e02,
            4.50708890e02,
            4.33088502e02,
            4.15885138e02,
            3.99097246e02,
            3.82723078e02,
            3.66760783e02,
            3.51208311e02,
            3.36063513e02,
            3.21323986e02,
            3.06987184e02,
            2.93050454e02,
            2.79510897e02,
            2.66365512e02,
            2.53611100e02,
            2.41244361e02,
            2.29261844e02,
            2.17659797e02,
            2.06434425e02,
            1.95581772e02,
            1.85097641e02,
            1.74977731e02,
            1.65217644e02,
            1.55812775e02,
            1.46758379e02,
            1.38049500e02,
            1.29681042e02,
            1.21647854e02,
            1.13944585e02,
            1.06565734e02,
            9.95056250e01,
            9.27584798e01,
            8.63183864e01,
            8.01792901e01,
            7.43350047e01,
            6.87792048e01,
            6.35054499e01,
            5.85071744e01,
            5.37776926e01,
            4.93102089e01,
            4.50978073e01,
            4.11334673e01,
            3.74100678e01,
            3.39203935e01,
            3.06571276e01,
            2.76128649e01,
            2.47801240e01,
            2.21513493e01,
            1.97189141e01,
            1.74751173e01,
            1.54122078e01,
            1.35223946e01,
            1.17978408e01,
            1.02306666e01,
            8.81297618e00,
            7.53686238e00,
            6.39440891e00,
            5.37770790e00,
            4.47887025e00,
            3.69003722e00,
            3.00339331e00,
            2.41118145e00,
            1.90571877e00,
            1.47941371e00,
            1.12478403e00,
            8.34477191e-01,
            6.01292929e-01,
            4.18208541e-01,
            2.78406802e-01,
            1.75308171e-01,
            1.02608084e-01,
            5.43204621e-02,
        ]
    )
    * 100
)
TEMP_GRID = np.arange(190, 311, 10)


class HITRANDatabase(CachedDatabase, OpticalDatabaseGenericAbsorber):
    def __init__(
        self,
        molecule: str,
        start_wavenumber: float,
        end_wavenumber: float,
        wavenumber_resolution: float,
        reduction_factor: int,
        num_threads: int = 1,
        db_root: Path | None = None,
        backend: str = "sasktran2",
        profile: str = "voigt",
    ) -> None:
        """


        Parameters
        ----------
        molecule : str
            The HITRAN molecule identifier, e.g. "CO2", "H2O", "O3", "CH4"
        start_wavenumber : float
            The start wavenumber of the database in cm^-1
        end_wavenumber : float
            The end wavenumber of the database in cm^-1
        wavenumber_resolution : float
            The wavenumber resolution of the database in cm^-1
        num_threads: int, optional
            The number of threads to use in the cross section calculation. Default is 1.
        db_root : Path, optional
            The root directory to store the database, by default None
        backend : str, optional
            The backend to use, by default "sasktran2" which uses the internal broadening algorithms.
            options are ["sasktran2", "sasktran_legacy", "hapi"]. Default is "sasktran2".
            Note: "sasktran_legacy" requires the sasktran package to be installed.
            Both the "sasktran2" and "hapi" backends use the hapi package to access the HITRAN database.
        profile : str, optional
            The line shape profile to use in the cross section calculation. Currently, only the "hapi" backend supports options other
            than the default "voigt". Supported options are ["voigt", "sdvoigt", "ht", "priority", "lorentz", "doppler"].
            Default is "voigt".
        """

        class NumpyEncoder(json.JSONEncoder):
            def default(self, obj):
                if isinstance(obj, np.ndarray):
                    return obj.tolist()
                return json.JSONEncoder.default(self, obj)

        self._profile = profile.lower()
        if backend == "sasktran_legacy" and self._profile != "voigt":
            logging.warning(
                "%s line profile is not supported by sasktran_legacy backend. Switching to voigt",
                self._profile,
            )
            self._profile = "voigt"

        hasher = hashlib.sha1()
        encoded = json.dumps(
            {
                "start_wavenumber": start_wavenumber,
                "end_wavenumber": end_wavenumber,
                "wavenumber_resolution": wavenumber_resolution,
                "reduction_factor": reduction_factor,
                "profile": self._profile,
            },
            sort_keys=True,
            cls=NumpyEncoder,
        ).encode()
        hasher.update(encoded)
        identifier = hasher.hexdigest()
        self._backend = backend
        self._start_wavenumber = start_wavenumber
        self._end_wavenumber = end_wavenumber
        self._wavenumber_resolution = wavenumber_resolution
        self._reduction_factor = reduction_factor
        self._molecule = molecule
        self._ltol = 1e-9
        self._wavenumber_wing = 50
        self._wavenumber_margin = 10
        self._num_threads = num_threads

        CachedDatabase.__init__(
            self,
            db_root=db_root,
            rel_path=Path("hitran").joinpath(self._molecule, self._backend),
        )

        self._data_file = self._db_root.joinpath(f"{identifier}.nc")

        OpticalDatabaseGenericAbsorber.__init__(self, self.path())

    def generate(self):
        if self._backend == "sasktran2":
            self._generate_sasktran2()
        elif self._backend == "sasktran_legacy":
            self._generate_sasktran_legacy()
        elif self._backend == "hapi":
            self._generate_hapi()
        else:
            msg = f"Invalid backend {self._backend}"
            raise ValueError(msg)

    def clear(self):
        if self._data_file.exists():
            self._data_file.unlink()

    def path(self, key: str | None = None, **kwargs) -> Path | None:  # noqa: ARG002
        if not self._data_file.exists():
            self.generate()
        return self._data_file

    def load_ds(self, key: str | None = None, **kwargs) -> xr.Dataset:
        return xr.open_dataset(self.path(key, **kwargs))

    def _generate_sasktran_legacy(self):
        try:
            from sasktran import ClimatologyUserDefined, HITRANChemical
        except ImportError as err:
            msg = "sasktran_legacy is required to generate Mie databases, try pip install sasktran"
            raise ImportError(msg) from err

        hires_wavenumber_grid = np.arange(
            self._start_wavenumber, self._end_wavenumber, self._wavenumber_resolution
        )

        xs = np.zeros((len(PRESSURE_GRID), len(TEMP_GRID), len(hires_wavenumber_grid)))

        for idx, pres in enumerate(PRESSURE_GRID):
            for idy, temp in enumerate(TEMP_GRID):
                xs[idx, idy] = (
                    HITRANChemical(self._molecule, line_tolerance=self._ltol)
                    .calculate_cross_sections(
                        ClimatologyUserDefined(
                            altitudes=[0, 100000],
                            values={
                                "SKCLIMATOLOGY_PRESSURE_PA": [pres, pres],
                                "SKCLIMATOLOGY_TEMPERATURE_K": [temp, temp],
                            },
                        ),
                        0,
                        0,
                        10000,
                        54372,
                        wavenumber_cminv_to_wavlength_nm(hires_wavenumber_grid),
                    )
                    .total
                    / 1e4
                )

        ds = xr.Dataset(
            {"xs": (["pressure", "temperature", "wavenumber_cminv"], xs)},
            coords={
                "pressure": PRESSURE_GRID,
                "temperature": TEMP_GRID,
                "wavenumber_cminv": hires_wavenumber_grid,
            },
        )

        ds.to_netcdf(self._data_file)

    def _generate_hapi(self):
        line_db = HITRANLineDatabase()
        line_db.path(self._molecule)

        hapi = get_hapi()

        abs_cal_fn = {
            "voigt": hapi.absorptionCoefficient_Voigt,
            "sdvoigt": hapi.absorptionCoefficient_SDVoigt,
            "ht": hapi.absorptionCoefficient_HT,
            "priority": hapi.absorptionCoefficient_Priority,
            "lorentz": hapi.absorptionCoefficient_Lorentz,
            "doppler": hapi.absorptionCoefficient_Doppler,
        }
        if self._profile not in abs_cal_fn:
            msg = f"Invalid line profile {self._profile}"
            raise ValueError(msg)

        line_db.initialize_hapi(self._molecule)

        hapi.select(
            self._molecule,
            DestinationTableName="spectral_window",
            Conditions=(
                "between",
                "nu",
                self._start_wavenumber - self._wavenumber_margin,
                self._end_wavenumber + self._wavenumber_margin,
            ),
        )

        hires_wavenumber_grid = np.arange(
            self._start_wavenumber, self._end_wavenumber, self._wavenumber_resolution
        )

        xs = np.zeros((len(PRESSURE_GRID), len(TEMP_GRID), len(hires_wavenumber_grid)))

        for idx, pres in enumerate(PRESSURE_GRID):
            for idy, temp in enumerate(TEMP_GRID):
                _, xs_hapi = abs_cal_fn[self._profile](
                    SourceTables="spectral_window",
                    Environment={"T": temp, "p": pres / 101325.0},
                    WavenumberGrid=hires_wavenumber_grid.tolist(),
                    WavenumberWing=self._wavenumber_wing,
                )
                xs[idx, idy] = xs_hapi / 1e4

        ds = xr.Dataset(
            {"xs": (["pressure", "temperature", "wavenumber_cminv"], xs)},
            coords={
                "pressure": PRESSURE_GRID,
                "temperature": TEMP_GRID,
                "wavenumber_cminv": hires_wavenumber_grid,
            },
        )

        ds.to_netcdf(self._data_file)

    def _generate_sasktran2(self):
        hires_wavenumber_grid = np.arange(
            self._start_wavenumber, self._end_wavenumber, self._wavenumber_resolution
        )

        calculator = HITRANAbsorber(
            self._molecule, line_contribution_width=self._wavenumber_wing / 2
        )
        xs = np.zeros((len(PRESSURE_GRID), len(TEMP_GRID), len(hires_wavenumber_grid)))

        for i, temperature in enumerate(TEMP_GRID):
            xs[:, i, :] = calculator.cross_sections(
                temperature_k=np.ones(len(PRESSURE_GRID)) * temperature,
                pressure_pa=PRESSURE_GRID,
                wavelengths_nm=wavenumber_cminv_to_wavlength_nm(hires_wavenumber_grid),
                num_threads=self._num_threads,
                altitudes_m=np.ones(len(PRESSURE_GRID)) * 0,
            ).T
        ds = xr.Dataset(
            {"xs": (["pressure", "temperature", "wavenumber_cminv"], xs)},
            coords={
                "pressure": PRESSURE_GRID,
                "temperature": TEMP_GRID,
                "wavenumber_cminv": hires_wavenumber_grid,
            },
        )

        ds.to_netcdf(self._data_file)
