import copy
import os
from datetime import datetime

import geopandas as gpd
import numpy as np
import pandas as pd
import pyproj
from scipy.interpolate import CubicSpline, interp1d
from shapely.geometry import LineString, Point
from shapely.ops import unary_union

from .spiderweb import TropicalCycloneSpiderweb
from .utils import gdf_to_geojson_js, gdf_to_pli

# Settings
dateformat_module = "%Y%m%d %H%M%S"
knots_to_ms = float(0.51444)
nm_to_km = float(1.852)
nm_to_m = float(1.852) * 1000

geodesic = pyproj.Geod(ellps="WGS84")

class TropicalCycloneEnsemble:
    # Init
    def __init__(self, tropical_cyclone,
                 name="ensemble",
                 number_of_realizations=10,
                 compute_wind_fields=True,
                 dt=3,
                 tstart=None,
                 tend=None,
                 duration=None,
                 track_path=None,
                 spw_path=None,
                 mean_abs_cte24=19.0397,
                 sc_cte=1.3253,
                 mean_abs_ate24=26.224,
                 sc_ate=1.3432,
                 mean_abs_ve24=6.9858,
                 sc_ve=1.0000,
                 bias_ve=0.0,                    
                 format="ascii"):

        """A tropical cyclone ensemble contains a collection (list) of tropical cyclone tracks that are generated by perturbing the best track of a tropical cyclone."""

        # If name is not given, use the name of the tropical cyclone
        self.name = name

        self.compute_wind_fields = compute_wind_fields

        if track_path is None:
            track_path = os.getcwd()
        if spw_path is None:
            spw_path = os.getcwd()

        self.track_path = track_path    
        if not os.path.exists(self.track_path):
            os.makedirs(self.track_path)

        self.spw_path = spw_path    
        if not os.path.exists(self.spw_path):
            os.makedirs(self.spw_path)

        self.format = format    

        self.number_of_realizations = number_of_realizations
        self.dt = dt

        # Error statistics - defaults are based on NHC of 2018-2021
        self.mean_abs_cte24 = mean_abs_cte24 * nm_to_m # mean absolute error in cross-track error (CTE) in meter
        self.sc_cte = sc_cte  # auto-regression CTE; typically >1
        self.mean_abs_ate24 = mean_abs_ate24 * nm_to_m # mean absolute error in along-track error (ATE) in meter
        self.sc_ate = sc_ate  # auto-regression ATE; typically >1
        self.mean_abs_ve24 = mean_abs_ve24  # mean absolute error in wind error (VE) in knots
        self.sc_ve = sc_ve  # auto-regression VE = 1 = no auto-regression
        self.bias_ve = bias_ve  # bias per hour

        self.tropical_cyclone = copy.deepcopy(tropical_cyclone)

        # Set scale factor of the best track to 1.0
        self.tropical_cyclone.track.gdf.loc[:,"wind_scale_factor"] = pd.Series(np.zeros(len(self.tropical_cyclone.track.gdf)) + 1.0,
                                                                               index=self.tropical_cyclone.track.gdf.index)

        if tstart is None:
            self.tstart = datetime.strptime(
                self.tropical_cyclone.track.gdf.datetime[0], dateformat_module
            )
        else:
            self.tstart = tstart

        if tend is not None:
            # Do we really want to shorten the track?
            self.tropical_cyclone.track.shorten(tend=tend)
        else:
            if duration is not None:
                tend = self.tstart + pd.Timedelta(duration, unit="h")
                self.tropical_cyclone.track.shorten(tend=tend)

        # Make sure the metric track is computed
        self.tropical_cyclone.compute_metric_track()

        # The actual ensemble members
        self.members = []
        member = TropicalCycloneEnsembleMember()
        member.tropical_cyclone = self.tropical_cyclone
        self.members.append(member)

    def generate(self):
        """Compute the ensemble members."""

        # Create the ensemble members 
        for i in range(self.number_of_realizations):
            member = TropicalCycloneEnsembleMember()
            # member.name = self.name + "_member_" + str(i)
            # Copy the track from the best track but not the spiderweb wind field
            member.tropical_cyclone = copy.deepcopy(self.tropical_cyclone)
            # But remove the spiderweb wind field to save memory
            member.tropical_cyclone.spiderweb = TropicalCycloneSpiderweb()
            self.members.append(member)

        self.generate_tracks() # Generates the tracks and vmax values for the individual ensemble members

        if self.compute_wind_fields:
            self.compute_wind_fields() # Computes the wind fields by scaling

    def generate_tracks(self):

        # Create new track with constant time intervals
        # Resample the track to hourly values
        equidistant_best_track = copy.copy(self.tropical_cyclone.track_metric)
        equidistant_best_track.resample(self.dt)

        for i in range(0, self.number_of_realizations):

            print("Generating ensemble track " + str(i + 1) + " of " + str(self.number_of_realizations) + " ...")

            member = self.members[i]
            member.equidistant_best_track = equidistant_best_track

            if i>0:
                member.generate_random_track(self.tstart,
                                            self.mean_abs_ate24,
                                            self.sc_ate,
                                            self.mean_abs_cte24,
                                            self.sc_cte,
                                            self.mean_abs_ve24,
                                            self.sc_ve,
                                            self.bias_ve)

            fname = os.path.join(self.track_path, self.name + str(i).zfill(5) + ".cyc")
            member.tropical_cyclone.track.write(fname)

    def compute_wind_fields(self, option="scaling"):

        for i in range(self.number_of_realizations):

            print("Generating ensemble wind field " + str(i + 1) + " of " + str(self.number_of_realizations) + " ...")

            if option == "scaling":

                member = self.members[i]

                # Wind field of best track
                spw = copy.deepcopy(self.tropical_cyclone.spiderweb)

                phi = spw.ds["azimuth"].values
                r = spw.ds["range"].values

                # Loop over time, set eye coordinates, and scale wind and pressure field 
                for it in range(len(member.tropical_cyclone.track.gdf)):
                    x0 = member.tropical_cyclone.track.gdf.geometry[it].coords[0][0]
                    y0 = member.tropical_cyclone.track.gdf.geometry[it].coords[0][1]
                    spw.ds["longitude_eye"].values[it] = x0
                    spw.ds["latitude_eye"].values[it] = y0
                    facw = member.tropical_cyclone.track.gdf.wind_scale_factor[it]                    
                    spw.ds["wind_x"].values[it, :, :] = spw.ds["wind_x"].values[it, :, :] * facw
                    spw.ds["wind_y"].values[it, :, :] = spw.ds["wind_y"].values[it, :, :] * facw
                    # Scale pressure with Vatvani relation
                    pavg = 100 * member.tropical_cyclone.config["background_pressure"]
                    spw.ds["pressure"].values[it, :, :] = pavg - (pavg - spw.ds["pressure"].values[it, :, :]) * facw**2
                    # Check if there is a dataarray for precipitation
                    if "precipitation" in spw.ds:
                        spw.ds["precipitation"].values[it, :, :] = spw.ds["precipitation"].values[it, :, :] * facw

                    # Create circular grid
                    pp, rr = np.meshgrid(phi, r)
                    xx = rr * np.cos(pp * np.pi / 180)
                    yy = rr * np.sin(pp * np.pi / 180)
                    # Convert to degrees
                    spw.ds["lon"].values[it, :, :] = x0 + xx / (111320 * np.cos(y0 * np.pi / 180)) 
                    spw.ds["lon"].values[it, :, :] = y0 + yy / 111320


                if self.format == "ascii":
                    ext = ".spw"
                else:
                    ext = ".nc"
                fname = os.path.join(self.spw_path, self.name + str(i).zfill(5) + ext)
                spw.write(fname,
                          format=self.format,
                          background_pressure=1013.0,
                          include_rainfall=True,
                          merge_frac=0.5)    

    def to_gdf(self,
               option="tracks",
               varname="track_ensemble",
               filename=None,
               buffer=300000.0,
               only_forecast=True):

        """Return the ensemble as a GeoDataFrame. Option can be either tracks or outline. varname is used when the format is js."""

        # Get the GDF
        if option == "tracks":
            gdf = self.tracks_to_gdf()
        elif option == "outline" or option == "cone":
            gdf = self.outline_to_gdf(buffer, only_forecast)

        # Write to file 
        if filename is not None:
            # Get extension of filename
            ext = os.path.splitext(filename)[-1]
            # Get path of filename
            folder_path = os.path.dirname(filename)
            # Make path (if needed)
            if folder_path != "":
                os.makedirs(folder_path, exist_ok=True)
            if ext == ".shp":
                gdf.to_file(filename)
            elif ext == ".geojson":
                gdf.to_file(filename, driver="GeoJSON")
            elif ext == ".js":
                gdf_to_geojson_js(gdf, filename, varname=varname)
            elif ext == ".pli":
                gdf_to_pli(gdf, filename)

        return gdf    

    def tracks_to_gdf(self):
        """Return the ensemble as a GeoDataFrame."""
        geometries = []
        ids = []
        for nn, member in enumerate(self.members):
            coordinates = []
            for it in range(len(member.tropical_cyclone.track.gdf)):
                coordinates.append(
                    (member.tropical_cyclone.track.gdf.geometry[it].x, member.tropical_cyclone.track.gdf.geometry[it].y)
                )
            line = LineString(coordinates)
            geometries.append(line)
            ids.append(member.name)
        gdf = gpd.GeoDataFrame({"geometry": geometries, "id": ids})
        return gdf

    def outline_to_gdf(self, buffer, only_forecast=True):
        """Generate GeoDataFrame with outline polygon of all ensemble members"""
        pols = []
        for member in self.members:
            # Loop through all points to make linestring
            coordinates = []
            for it in range(len(member.tropical_cyclone.track.gdf)):
                t = datetime.strptime(member.tropical_cyclone.track.gdf.datetime[it], dateformat_module)
                if only_forecast and t < self.tstart:
                    # This is part of the best track
                    continue
                coordinates.append(
                    (member.tropical_cyclone.track.gdf.geometry[it].x, member.tropical_cyclone.track.gdf.geometry[it].y)
                )
            line = LineString(coordinates)
            line = (
                gpd.GeoDataFrame({"geometry": [line]})
                .set_crs(4326)
                .to_crs(3857)
                .iloc[0]["geometry"]
            )
            pol = line.buffer(buffer).simplify(0.1)
            pols.append(pol)
        outline_geom = unary_union(pols)
        gdf = gpd.GeoDataFrame({"geometry": [outline_geom]}).set_crs(3857).to_crs(4326)
        return gdf

    # # Make visual of all the tracks
    # def make_figures(self, folder_path):
    #     # Make path (if needed)
    #     os.makedirs(folder_path, exist_ok=True)

    #     # Get times
    #     datetime_ensemble = []
    #     for it in range(len(self.members[0].track)):
    #         datetime_ensemble.append(
    #             datetime.strptime(self.members[0].track.datetime[it], dateformat_module)
    #         )

    #     # Show wind speeds
    #     for nn, member in enumerate(self.members):
    #         plt.plot(datetime_ensemble, member.track.vmax)
    #     plt.show()

class TropicalCycloneEnsembleMember:

    def __init__(self):
        self.name = None
        self.tropical_cyclone = None

    def generate_random_track(self,
                              tstart_ensemble,
                              mean_abs_ate24,
                              sc_ate,
                              mean_abs_cte24,
                              sc_cte,
                              mean_abs_ve24,
                              sc_ve,
                              bias_ve):

        # Generate random track

        # Number of times in resampled best track
        ntpred = len(self.equidistant_best_track.gdf)

        # Get time step
        t0 = datetime.strptime(self.equidistant_best_track.gdf.datetime[0], dateformat_module)
        t1 = datetime.strptime(self.equidistant_best_track.gdf.datetime[1], dateformat_module)
        dt_seconds = (t1 - t0).total_seconds()
        dtd = dt_seconds / 86400

        np.random.seed()

        # Random error matrices from normal distribution
        arnd0 = np.random.randn(ntpred)
        crnd0 = np.random.randn(ntpred)
        vrnd0 = np.random.randn(ntpred)

        # Limit to -2 and +2 sigma
        arnd0 = np.maximum(np.minimum(arnd0, 2.0), -2.0)
        crnd0 = np.maximum(np.minimum(crnd0, 2.0), -2.0)
        vrnd0 = np.maximum(np.minimum(vrnd0, 2.0), -2.0)

        # Rest
        ate  = np.zeros((ntpred))
        cte  = np.zeros((ntpred))
        ve   = np.zeros((ntpred))
        x    = np.zeros((ntpred))
        y    = np.zeros((ntpred))
        vmax = np.zeros((ntpred))
        wind_scale = np.zeros((ntpred)) + 1.0

        # Correct for time
        tfac = np.sqrt(dtd)

        # go from mean absolute error to standard deviation
        sigma_ate = mean_abs_ate24 / np.sqrt(2 / np.pi)
        sigma_cte = mean_abs_cte24 / np.sqrt(2 / np.pi)

        # Check unit of intensity
        if self.equidistant_best_track.unit_intensity == "knots":
            sigma_ve = mean_abs_ve24 / np.sqrt(2 / np.pi)
        else:
            sigma_ve = mean_abs_ve24 * knots_to_ms / np.sqrt(2 / np.pi)

        # Loop over time
        forecast_timestep = 0.0

        # Find index of last time smaller or equal to tstart_ensemble
        it_start = 0
        for it in range(0, len(self.equidistant_best_track.gdf)):
            t = datetime.strptime(self.equidistant_best_track.gdf.datetime[it], dateformat_module)
            if t <= tstart_ensemble:
                it_start = it
                break

        # Coordinates of the first point
        x00 = self.equidistant_best_track.gdf.geometry[it_start].x
        y00 = self.equidistant_best_track.gdf.geometry[it_start].y

        # Initial errors 
        ate_12 = 0.0
        cte_12 = 0.0
        ve_12  = 0.0    

        # Loop through times in equidistant best track gdf
        for it in range(ntpred):

            t = datetime.strptime(self.equidistant_best_track.gdf.datetime[it], dateformat_module)
            x0 = self.equidistant_best_track.gdf.geometry[it].x
            y0 = self.equidistant_best_track.gdf.geometry[it].y
            vmax0 = self.equidistant_best_track.gdf.vmax[it]       

            x[it] = x0
            y[it] = y0
            vmax[it] = vmax0

            forecast_timestep = forecast_timestep + dtd * 24

            # If we want the variability
            if t > tstart_ensemble:

                # standard deviation scales with tfac
                arnd = tfac * sigma_ate * arnd0[it]
                crnd = tfac * sigma_cte * crnd0[it]
                vrnd = tfac * sigma_ve * vrnd0[it]

                # Limit to -2 and +2 sigma
                arnd = np.maximum(
                    np.minimum(arnd, 2 * tfac * sigma_ate), -2 * tfac * sigma_ate
                )
                crnd = np.maximum(
                    np.minimum(crnd, 2 * tfac * sigma_cte), -2 * tfac * sigma_cte
                )
                vrnd = np.maximum(
                    np.minimum(vrnd, 2 * tfac * sigma_ve), -2 * tfac * sigma_ve
                )

                # Compute new track errors
                at1 = sc_ate**dtd
                ate[it] = at1 * ate_12 + arnd
                ct1 = sc_cte**dtd
                cte[it] = ct1 * cte_12 + crnd

                # Compute track heading based on latitude / longitude of two points
                fwd_azimuth, back_azimuth, distance = geodesic.inv(x00, y00, x0, y0)

                # Along track error shift
                xa, ya, backaz = geodesic.fwd(
                    x0,
                    y0,
                    fwd_azimuth,
                    ate[it],
                    radians=False,
                )

                # Cross track error shift
                xc, yc, backaz = geodesic.fwd(
                    xa, ya, fwd_azimuth + 90.0, cte[it], radians=False
                )

                x[it] = xc
                y[it] = yc

                # Do wind speed too
                ve1 = sc_ve**dtd
                ve[it] = ve1 * ve_12 + vrnd + bias_ve * forecast_timestep
                # Let's make sure that the wind speed does not become negative
                if vmax0 + ve[it] < 5.0:
                    ve[it] = max(5.0 - vmax0, 0.0)
                vmax[it] = vmax0 + ve[it]
                wind_scale[it] = vmax[it] / vmax0

                # Save errors from last iteration
                ate_12 = ate[it]
                cte_12 = cte[it]
                ve_12  = ve[it]

        # And now interpolate x, y, vmax and wind_scale on the equidistant best track, to the original time steps

        # First make time stamp arrays for equidistant and original best track
        t0 = []
        for date in self.equidistant_best_track.gdf["datetime"]:
            t0.append(datetime.strptime(date, dateformat_module))
        t1 = []
        for date in self.tropical_cyclone.track.gdf["datetime"]:
            t1.append(datetime.strptime(date, dateformat_module))
        t0 = [date.timestamp() for date in t0]
        t1 = [date.timestamp() for date in t1]

        # Now interpolate

        # Use CubicSpline for x and y
        f = CubicSpline(t0, x)
        x1 = f(t1)
        f = CubicSpline(t0, y)
        y1 = f(t1)

        # Linear interpolation for vmax and wind_scale
        f = interp1d(t0, vmax)
        v1 = f(t1) / knots_to_ms
        f = interp1d(t0, wind_scale)
        ws1 = f(t1)

        # Also update central pressure 
        pavg = self.tropical_cyclone.config["background_pressure"]
        pc1 = np.zeros(len(t1))
        for it in range(len(t1)):
            pc1[it] = pavg - (pavg - self.tropical_cyclone.track.gdf.pc[it]) * ws1[it]**2

        # Loop over time for geometry
        points = []
        for it in range(len(t1)):
            points.append(Point(x1[it], y1[it]))
            self.tropical_cyclone.track.gdf.loc[it, "geometry"] = Point(x1[it], y1[it])

        # And set Vmax and wind scale factor
        self.tropical_cyclone.track.gdf.loc[:,"vmax"] = pd.Series(v1, index=self.tropical_cyclone.track.gdf.index)
        self.tropical_cyclone.track.gdf.loc[:,"pc"] = pd.Series(pc1, index=self.tropical_cyclone.track.gdf.index)
        self.tropical_cyclone.track.gdf.loc[:,"wind_scale_factor"] = pd.Series(ws1, index=self.tropical_cyclone.track.gdf.index)

        # Recompute metric values
        self.tropical_cyclone.compute_metric_track()
