import os, netCDF4, cfgrib
import numpy as np
import xarray as xr
import datetime

this_dir = os.path.dirname(__file__)
NC_FILE =  os.path.join(this_dir, "sample.nc")
GRIB_FILE = os.path.join(this_dir, "sample.grib")


def pytest_addoption(parser):
    parser.addoption("--ncfile", action="store", default=NC_FILE,
                     help="netCDF4 file")
    parser.addoption("--gribfile", action="store", default=GRIB_FILE,
                     help="grib file")


def pytest_configure():
    if not os.path.exists(NC_FILE):
        print("Creating fake datasets (netcdf and grib formats)")
        mock_nc()
    if not os.path.exists(GRIB_FILE):
        mock_grib()


def pytest_unconfigure():
    print("Removing fake datasets")
    if os.path.exists(NC_FILE):
        os.remove(NC_FILE)
    if os.path.exists(GRIB_FILE):
        os.remove(GRIB_FILE)


def mock_nc():
    ds = None
    try:
        # Mocked dimension/variables
        lon = np.linspace(0, 360, 320, endpoint=False)
        lat = np.linspace(-90, 90, 160, endpoint=False)
        tas = np.random.randint(10, 25,
                              size=(1, len(lon), len(lat)))
        fake_date = datetime.date.today()
        fake_dates = [fake_date + datetime.timedelta(days=i) for i in range(5)]
        num_dates = np.array([d.toordinal() for d in fake_dates])

        # Writing netcDF4 and adding mocks
        ds = netCDF4.Dataset(NC_FILE, 'w', format='NETCDF4')
        ds.set_ncstring_attrs(True)
        ds.createDimension('lon', len(lon))
        ds.createDimension('lat', len(lat))
        ds.createDimension('time', None)
        ds.createDimension('tas', len(tas))
        longitude = ds.createVariable('lon', 'f4', 'lon')
        latitude = ds.createVariable('lat', 'f4', 'lat')
        time = ds.createVariable('time', 'f4', 'time')
        temp = ds.createVariable('tas', 'f4', ('time', 'lon', 'lat'))
        longitude[:] = lon
        latitude[:] = lat
        temp[:] = tas
        time[:5] = num_dates

        # Adding attributes
        ds.name = "Example dataset"
        ds.description = "Example dataset"
        longitude.units = 'degrees east'
        latitude.units = 'degrees north'
        time.units = 'days since 1-01-01'
        temp.units = 'Kelvin'
        temp.warning = 'This data is not real!'
    finally:
        ds.close()


def mock_grib():
   ds = xr.open_dataset(NC_FILE)
   ds = ds.isel(time=0)
   cfgrib.to_grib(ds, GRIB_FILE, grib_keys={"gridType": "unstructured_grid"})


if __name__ == "__main__":
    mock_nc()
    mock_grib()
