#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import netCDF4 as nc
import settings


def generate_basic_ncfile(outfile, crs=None):
    
    """
        Initialise basic netCDF4 file
        
        Input:
            Outfile: path to store the netcdf file, string
            crs: coordinate reference system, string
    """

    # If outfile exists already, delete
    if os.path.exists(outfile):
        os.remove(outfile)

    ds = nc.Dataset(outfile, 'w', format='NETCDF4')

    return ds


def generate_ncfile(outfile, x, y, data, crs=None,
                    data_unit=None, missing_value=settings.no_value):

    """
        Save 2D dataset as netCDF4 file
        
        Input:
            Outfile: path to store the netcdf file, string
            x: longitude vector, list
            y: latitude vector, list
            data: 2D data array
            crs: coordinate reference system, string
            data_unit: data unit, string
            missing_value: no data value, integer or float
    """

    # If outfile exists already, delete
    if os.path.exists(outfile):
        os.remove(outfile)

    ds = nc.Dataset(outfile, 'w', format='NETCDF4')
    ds.createDimension('lon', len(x))
    ds.createDimension('lat', len(y))
    longitude = ds.createVariable('Longitude', 'f4', 'lon')
    latitude = ds.createVariable('Latitude', 'f4', 'lat')
    result = ds.createVariable('Result', 'f4', ('lat', 'lon'))

    longitude[:] = x
    latitude[:] = y
    result[:, :] = data

    # Provide global information in output-file
    if crs is not None:
        longitude.units = crs
        latitude.units = crs
    if data_unit is not None:
        result.units = data_unit
    ds.missing_value = missing_value
    ds.close()


def generate_3dncfile(outfile, x, y, data, dim, features, crs='wgs84',
                      data_unit=None, missing_value=settings.no_value):
    
    """
        Save 3D dataset as netCDF4 file, e.g. data cube
        
        Input:
            Outfile: path to store the netcdf file, string
            x: longitude vector, list
            y: latitude vector, list
            dim: number of 2D datasets, integer
            data: 2D data array
            features: contained features in prediction dataset, list of chars
            crs: coordinate reference system, string
            data_unit: data unit, string
            missing_value: no data value, integer or float
    """

    # If outfile exists already, delete
    if os.path.exists(outfile):
        os.remove(outfile)

    ds = nc.Dataset(outfile, 'w', format='NETCDF4')
    ds.createDimension('lon', len(x))
    ds.createDimension('lat', len(y))
    ds.createDimension('dim', dim)
    ds.createDimension('feat', len(features))
    longitude = ds.createVariable('Longitude', 'f4', 'lon')
    latitude = ds.createVariable('Latitude', 'f4', 'lat')
    result = ds.createVariable('Result', 'f4', ('lat', 'lon', 'dim'))
    Features = ds.createVariable('features', 'S1', 'feat')

    longitude[:] = x
    latitude[:] = y
    result[:, :, :] = data
    Features[:] = features

    # Provide global information in output-file
    if crs is not None:
        longitude.units = crs
        latitude.units = crs
    if data_unit is not None:
        result.units = data_unit
    ds.missing_value = missing_value
    ds.close()


def generate_2dncfile(outfile, x, y, data, features, crs='wgs84',
                      data_unit=None, missing_value=settings.no_value):
    
    """
        Save 2D dataset as netCDF4 file, e.g. Prediction dataset
        
        Input:
            Outfile: path to store the netcdf file, string
            x: longitude vector, list
            y: latitude vector, list
            data: 2D data array
            features: contained features in prediction dataset, list of chars
            crs: coordinate reference system, string
            data_unit: data unit, string
            missing_value: no data value, integer or float
    """

    # If outfile exists already, delete
    if os.path.exists(outfile):
        os.remove(outfile)

    ds = nc.Dataset(outfile, 'w', format='NETCDF4')
    ds.createDimension('lon', len(x))
    ds.createDimension('lat', len(y))
    ds.createDimension('feat', len(features))
    longitude = ds.createVariable('Longitude', 'f4', 'lon')
    latitude = ds.createVariable('Latitude', 'f4', 'lat')
    result = ds.createVariable('Result', 'f4', ('lat', 'lon'))
    Features = ds.createVariable('features', 'S1', 'feat')

    longitude[:] = x
    latitude[:] = y
    result[:, :] = data
    Features[:] = features

    # Provide global information in output-file
    if crs is not None:
        longitude.units = crs
        latitude.units = crs
    if data_unit is not None:
        result.units = data_unit
    ds.missing_value = missing_value
    ds.close()