# from copy import deepcopy
from netCDF4 import Dataset
from mpi4py import MPI
import numpy as np
# from scipy.ndimage import zoom

CHIMERE_GLOBAL_ATTRS = {
    "Title": "CHIMERE SUITE",
    "Sub-title": "Emissions file - Surfacic + Point sources + Fires",
    "Generating_process": "Generated by Hermes_delta",
    "Conventions": "",
    "Domain": "EU027n",
    "history": "",
}
CONSTANTS = {"SpStrLen": 23, "DateStrlen": 19}
# SOUTH_NORTH = 1230
# WEST_EAST = 2050


def to_netcdf_chimere(self, path, keep_open=True):
    # Open NetCDF
    if self.info:
        print("Rank {0:03d}: Creating {1}".format(self.rank, path))
    if self.size > 1:
        netcdf = Dataset(
            path,
            format="NETCDF4",
            mode="w",
            parallel=True,
            comm=self.comm,
            info=MPI.Info(),
        )
    else:
        netcdf = Dataset(path, format="NETCDF4", mode="w", parallel=False)
    if self.info:
        print("Rank {0:03d}: NetCDF ready to write".format(self.rank))

    _create_dimensions(self, netcdf)
    _create_and_convert_time(self, netcdf)
    netcdf.createVariable("species", "S1", ("Species", "SpStrLen"))
    _create_coordinates_dimensions(self, netcdf)


    _create_coordinates_variables(self, netcdf)
    _create_species_variable(self, netcdf)
    netcdf.close()


def _create_coordinates_dimensions(self, netcdf):
    netcdf.createDimension("south_north", len(self.lat["data"]))
    netcdf.createDimension("west_east", len(self.lon["data"]))


def _create_coordinates_variables(self, netcdf):
    lat_orig = np.array(self.lat["data"])
    lon_orig = np.array(self.lon["data"])

    lon_2d, lat_2d = np.meshgrid(lon_orig, lat_orig)
    lon = netcdf.createVariable("lon", "f4", ("south_north", "west_east"))
    lon.units = "degrees_north"
    lon.long_name = "Longitude"
    lon[:, :] = lon_2d

    lat = netcdf.createVariable("lat", "f4", ("south_north", "west_east"))
    lat.units = "degrees_east"
    lat.long_name = "Latitude"
    lat[:, :] = lat_2d

# def calculate_grid_Area

def _create_and_convert_time(self, netcdf):
    formatted_times = [t.strftime("%Y-%m-%d_%H:%M:%S") for t in self.time]
    times_array = np.array(
        [list(s.ljust(CONSTANTS["DateStrlen"])) for s in formatted_times], dtype="S1"
    )

    netcdf.createDimension("Time", None)
    netcdf.createDimension("DateStrLen", CONSTANTS["DateStrlen"])

    time_var = netcdf.createVariable("Times", "S1", ("Time", "DateStrLen"))
    time_var[:, :] = times_array

# spcies, spstrlen, bottom_top
def _create_dimensions(self, netcdf):
    netcdf.createDimension("SpStrLen", CONSTANTS["SpStrLen"])
    netcdf.createDimension("Species", len(self.variables.keys()))
    netcdf.createDimension("bottom_top", len(self.lev["data"]))   

def _create_species_variable(self, netcdf):
    pollutants = list(self.variables.keys())
    pollutants_uppercase = [x.upper() for x in pollutants] 


    for pollutant in pollutants_uppercase:
        pollutant_orig = np.array(self.variables[pollutant.lower()]["data"])

        pollutants_variable = netcdf.createVariable(str(pollutant), "f4", ("Time", "bottom_top", "south_north", "west_east"))
        pollutants_variable.units = "molecule/cm2/s"
        pollutants_variable.long_name = f"{pollutant} Emission"

        # write data
        pollutants_variable[:, :, :, :] = pollutant_orig

