import numpy as np
from pratmo import pratmo
import xarray as xr
import sasktranif.sasktranif as skif


class ECMWF_with_deltaT:
    def __init__(self, deltaT, lat, lon, mjd, heights):
        self.ecmwf = skif.ISKClimatology('OSIRIS_ECMWF')
        self.ecmwf.UpdateCache([lat, lon, 0.0, mjd])
        self.climate = skif.ISKClimatology('USERDEFINED_PROFILE')
        self.climate.SetProperty('DoLogInterpolation', 1)
        if heights[0] > 0:
            heights.insert(0, 0)
        temp = []
        pres = []
        air = []
        for height in heights:
            ok, t = self.ecmwf.GetParameter(skif.SKCLIMATOLOGY_TEMPERATURE_K, [lat, lon, height, mjd])
            temp.append(t+deltaT)
            ok, p = self.ecmwf.GetParameter(skif.SKCLIMATOLOGY_PRESSURE_PA, [lat, lon, height, mjd])
            pres.append(p)
            ok, a = self.ecmwf.GetParameter(skif.SKCLIMATOLOGY_AIRNUMBERDENSITY_CM3, [lat, lon, height, mjd])
            air.append(a)
        self.climate.SetProperty('Heights', heights)
        self.climate.SetPropertyUserDefined(skif.SKCLIMATOLOGY_TEMPERATURE_K, temp)
        self.climate.SetPropertyUserDefined(skif.SKCLIMATOLOGY_PRESSURE_PA, pres)
        self.climate.SetPropertyUserDefined(skif.SKCLIMATOLOGY_AIRNUMBERDENSITY_CM3, air)


class Pratmo:
    """A wrapper for the pratmo++ interface."""
    def __enter__(self):
        self.model = pratmo.Pratmo_BoxModel()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def run_boxmodel(self, lat, lon, mjd, heights=None):
        if heights:
            self.model.SetBoxHeights(heights)
        res = self.model.Solve(lat, lon, mjd)
        if res:
            raise Exception('PRATMO: Solve returned non-zero. %s altitudes did not converge.'%res)

    def getparam(self, param):
        if param not in self.model.GetVariableNames('ARMLJ'):
            raise Exception('PRATMO: GetSpecies provided invalid species: ' + param)
        res, data = self.model.GetSpecies(param)
        if not res:
            raise Exception('PRATMO: GetSpecies failed to return valid data for species ' + param)
        return data

    def add_lst(self, lst):
        self.model.AddCustomLocalSolarTime(lst)

    def getparams_asdataset(self, params):
        # Rearrange the Local Solar Times
        lsts = self.model.LocalSolarTimes()
        lsts = [lst if lst < 12 else lst - 24 for lst in lsts]
        lsts[-1] = 12

        data_vars = {param: (['LST', 'Altitude'], self.getparam(param)) for param in params}

        ds = xr.Dataset(data_vars=data_vars,
                        coords={'LST': lsts,
                                'Altitude': np.array(self.model.BoxHeights())})
        return ds.transpose()

    def seto3(self, alts, o3):
        # Set the ozone profile
        if alts[0]>0:
            alts.insert(0, 0)
            o3.insert(0, o3[0])
        climate = skif.ISKClimatology('USERDEFINED_PROFILE')
        climate.SetProperty('DoLogInterpolation', 1)
        climate.SetProperty('Heights', alts)
        climate.SetPropertyUserDefined(skif.SKCLIMATOLOGY_O3_VMR, o3)
        self.model.SetClimatology('O3', climate)

    def setalbedo(self, albedo):
        self.model.SetAlbedo(albedo)

    def useOSIRISECMWFbackground(self):
        # Use ECMWF instead of the default PRATMO background.
        if not self.model.SetAtmosphericStateClimatology(skif.ISKClimatology('OSIRIS_ECMWF')):
            raise Exception('ISKClimatology: Failed to set climatology to OSIRIS_ECMWF')

    def setbackground(self, clim):
        # Use custom climatology as background.
        if not self.model.SetAtmosphericStateClimatology(clim):
            raise Exception('ISKClimatology: Failed to set climatology to OSIRIS_ECMWF')