import os, cfgrib, zipfile, shutil, json
import numpy as np
import pandas as pd
import xarray as xr
import datetime
from pathlib import Path
from types import SimpleNamespace

this_dir = os.path.dirname(__file__)
dirs = SimpleNamespace(
            test_dir = this_dir,
            nc_file=os.path.join(this_dir, "sample20200101.nc"),
            grib_file=os.path.join(this_dir, "sample20200101.grib"),
            grib_cds = os.path.join(this_dir, "cds_sample.grib"),
            zip_file = os.path.join(this_dir, "sample20200101.zip"),
            log_dir=os.path.join("/tmp", "dqc_logs", "test"),
            log_file=os.path.join("/tmp", "dqc_logs", "test", "test.log"),
            mask_file=os.path.join(this_dir, "mask.nc"),
            hist_csv=os.path.join(this_dir, "samplehist.csv"),
            stats_csv=os.path.join(this_dir, "samplestats.csv"),
            constraints_json = os.path.join(this_dir, "constraints.json"),
            mapping_json = os.path.join(this_dir, "mapping.json")
        )

FAKE_DATE = datetime.date(2019, 1, 1)

def pytest_configure():
    Path(dirs.log_dir).mkdir(parents=True, exist_ok=True)
    print("Creating fake datasets (netcdf, grib and zip formats)")
    mock_nc()
    mock_grib()
    mock_zip()
    mock_jsons()
    mock_mask()
    mock_logs()

def pytest_unconfigure():
    print("Removing fake files")
    Path(dirs.nc_file).unlink()
    Path(dirs.grib_file).unlink()
    Path(dirs.zip_file).unlink()
    Path(dirs.constraints_json).unlink()
    Path(dirs.mapping_json).unlink()
    Path(dirs.mask_file).unlink()
    if Path(dirs.hist_csv).exists():
        Path(dirs.hist_csv).unlink()
    if Path(dirs.stats_csv).exists():
        Path(dirs.stats_csv).unlink()
    Path(dirs.log_file).unlink()
    shutil.rmtree(dirs.log_dir)


def mock_nc():
    ds = None
    try:
        lon = np.linspace(0, 360, 320, endpoint=False)
        lat = np.linspace(-90, 90, 160, endpoint=False)
        tas = np.random.randint(10, 25,
                              size=(5, len(lon), len(lat)))
        fake_dates = pd.date_range(FAKE_DATE, periods=5)
        # fake_dates = [FAKE_DATE + datetime.timedelta(days=i) for i in range(5)]
        da = xr.DataArray(tas, coords=[fake_dates, lon, lat],
                          dims=['time', 'lon', 'lat'])
        ds = da.to_dataset(name='tas')
        ds.to_netcdf(dirs.nc_file)
    finally:
        if ds:
            ds.close()


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

def mock_zip():
    zip = zipfile.ZipFile(dirs.zip_file, "w", zipfile.ZIP_DEFLATED)
    zip.write(dirs.nc_file, os.path.basename(dirs.nc_file))
    zip.write(dirs.grib_file, os.path.basename(dirs.grib_file))
    zip.close()

def mock_jsons():
    constraints = [{
        "variable": ["tas"],
        "year": [FAKE_DATE.year],
        "month": [FAKE_DATE.month],
        "day": [FAKE_DATE.day, FAKE_DATE.day+1]
    }]
    with open(Path(dirs.constraints_json), 'w') as cf:
        json.dump(constraints, cf)
    mapping = {'remap': {'variable': {'tas': '171129'}}}
    with open(Path(dirs.mapping_json), 'w') as mf:
        json.dump(mapping, mf)

def mock_mask():
    ds = xr.open_dataset(dirs.nc_file)
    ds['tas'][0, 0, 0] = np.nan  # nan first point, values otherwise
    # [0] because for some reason time must be a dimension on the mask
    mask2d = ds.isel(time=[0])['tas']
    mask2d[-1, -1] = 0  # False last point, values otherwise
    mask2d = mask2d.compute().astype(bool)
    mask2d = ~mask2d
    mask2d.to_netcdf(dirs.mask_file)

def mock_logs():
    Path(dirs.log_dir).mkdir(parents=True, exist_ok=True)
    with open(dirs.log_file, 'w') as f:
        f.write("Log test")


if __name__ == "__main__":
    pytest_configure()
