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
from dqc_wrapper.test06_valid_ranges import ValidRangesTest
import yaml

this_dir = os.path.dirname(__file__)
log_dir = os.path.join("/tmp", "dqc_logs", "test")


def get_masks_dir():
    if Path("/shared/cds_downloads/masks").exists():
        return "/shared/cds_downloads/masks"
    else:
        return "/home/suareze/dqc/masks"


dirs = SimpleNamespace(
            test_dir=this_dir,
            masks_dir=get_masks_dir(),
            nc_file=os.path.join(this_dir, "sample20200101.nc"),
            nc_masks_file=os.path.join(this_dir, "sample_for_masks.nc"),
            grib_file=os.path.join(this_dir, "sample20200101.grib"),
            grib_cds=os.path.join(this_dir, "cds_sample.grib"),
            grib_json=os.path.join(this_dir, "../grib/grib1.json"),
            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(log_dir, "testname-hist.csv"),
            stats_csv=os.path.join(log_dir, "testname-stats.csv"),
            constraints_json=os.path.join(this_dir, "constraints.json"),
            mapping_json=os.path.join(this_dir, "mapping.json"),
            form_json=os.path.join(this_dir, "form.json"),
            variables_yaml=os.path.join(this_dir, "variables.yaml"),
            test_image=os.path.join(this_dir, "../utils/img/stats_explained.png"),
            conf_test=os.path.join(this_dir, "../check_tests/conf_utils_unit_tests_.ini"),
            txt_file=os.path.join(this_dir, '../utils/texts/caption_histogram_reanalysis.txt')
        )

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_nc_for_masks()
    mock_grib()
    mock_zip()
    mock_jsons()
    mock_mask()
    mock_logs()
    mock_valid_ranges()
    mock_conf_file()


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.form_json).unlink()
    Path(dirs.variables_yaml).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.conf_test).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)
        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_nc_for_masks():
    ds = None
    try:
        lon = np.linspace(0, 360, 1030, endpoint=False)
        lat = np.linspace(-90, 90, 1008, endpoint=False)
        tas = np.random.randint(10, 25,
                              size=(5, len(lon), len(lat)))
        fake_dates = pd.date_range(FAKE_DATE, periods=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_masks_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 = {
               "force": {
                   "class": [
                       "c3"
                   ],
                   "database": "ecmwf",
                   "day": [
                       "01"
                   ],
                   "expect": [
                       "any"
                   ],
                   "expver": [
                       "1"
                   ],
                   "levtype": [
                       "pl"
                   ],
                   "method": [
                       "1"
                   ],
                   "number": [
                       "all"
                   ],
                   "stream": [
                       "msmm"
                   ],
                   "time": [
                       "00:00"
                   ]
               },
               "options": {
                   "selection_limit_ignore": [
                       "area",
                       "grid"
                   ],
                   "wants_dates": True
               },
               "remap": {
                   "originating_centre": {
                       "dwd": "edzw",
                       "ecmwf": "ecmf",
                       "meteo_france": "lfpw",
                       "ncep": "kwbc",
                       "ukmo": "egrr"
                   },
                   "product_type": {
                       "ensemble_mean": "em",
                       "hindcast_climate_mean": "hcmean",
                       "monthly_mean": "fcmean"
                   },
                   "time": {
                       "00:00": "00:00:00"
                   },
                   "variable": {
                       "geopotential": "129",
                       "specific_humidity": "133",
                       "temperature": "130",
                       "u_component_of_wind": "131",
                       "v_component_of_wind": "132",
                       'tas': '171129'
                   }
               },
               "rename": {
                   "leadtime_hour": "step",
                   "leadtime_month": "fcmonth",
                   "originating_centre": "origin",
                   "pressure_level": "levelist",
                   "product_type": "type",
                   "variable": "param"
               },
               "selection_limit": 10000,
               "selection_limit_ignore": [
                   "area",
                   "grid"
               ]}
    with open(Path(dirs.mapping_json), 'w') as mf:
        json.dump(mapping, mf)
    form = [{
        "css": "todo",
        "details": {
            "columns": 2,
            "id": 2,
            "labels": {
                "geopotential": "Geopotential",
                "specific_humidity": "Specific humidity",
                "temperature": "Temperature",
                "u_component_of_wind": "U-component of wind",
                "v_component_of_wind": "V-component of wind"
            },
            "values": [
                "geopotential",
                "specific_humidity",
                "temperature",
                "u_component_of_wind",
                "v_component_of_wind"
            ]
        },
        "help": "Please, consult the product user guide in the documentation section for more information on these variables.",
        "label": "Variable",
        "name": "variable",
        "required": True,
        "type": "StringListWidget"
    }]
    with open(dirs.form_json, 'w') as ff:
        json.dump(form, ff)
    variables = {'variables': {'Geopotential': {'units': 'm^2 s^-2'},
                               'Specific humidity': {'units': 'kg kg^-1'},
                               'Temperature': {'units': 'K'},
                               'U-component of wind': {'units': 'm s^-1'},
                               'V-component of wind': {'units': 'm s^-1'}}}
    with open(dirs.variables_yaml, 'w') as vy:
        yaml.dump(variables, vy)


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")


def mock_conf_file():
    # using existing configuration file but changing dirs (needs absolute paths)
    old_path = os.path.join(this_dir, "../check_tests/conf_utils_unit_tests.ini")
    with open(old_path, 'r') as iconf:
        conf = iconf.read()
        newconf = conf.replace("/tmp", dirs.test_dir)
        with open(dirs.conf_test, 'w') as oconf:
            oconf.write(newconf)


def mock_valid_ranges():
    vr = ValidRangesTest(dirs.log_dir, 'testname.ini', [dirs.nc_file], 'tas', '', '',
                         'fc', '',  dirs.test_dir, 'tas', 'degC')
    vr.check_valid_ranges()
