#!/usr/bin/env python
# coding: utf-8

import os
import matplotlib.pyplot as plt
import numpy as np

from glob import glob
from netCDF4 import Dataset

import c2sm
import smartcarb as sc


def tracer_mass_ts(tracer='CO2_BV'):
    """\
    Time series of CO2 mass in emissions and model domain.
    """
    root = '/scratch/snx1600/gkuhl/smartcarb/mother/2015010100_0_72/'

    inp_filenames = os.path.join(root, 'int2lm', 'output', 'lbfd??????????e.nc')
    print(inp_filenames)
    inp_filenames = sorted(glob(inp_filenames))

    out_filenames = os.path.join(root, 'cosmo', 'output', 'lffd??????????.nc')
    print(out_filenames)
    out_filenames = sorted(glob(out_filenames))

    n_files = len(inp_filenames)


    emissions = np.zeros(n_files)
    transport = np.zeros(n_files)


    # height levels
    const_filename = os.path.join(root, 'cosmo', 'output', 'lffd2015010100c.nc')
    with Dataset(const_filename) as nc:
        h = nc.variables['HHL'][:]


    #area = 1003746.5 # in m2
    area = 1235941.1 # in m2

    for i, (ifn, ofn) in enumerate(zip(inp_filenames, out_filenames)):

        print(os.path.basename(ifn), os.path.basename(ofn))

        with Dataset(ifn) as nc:
            emissions[i+1] = emissions[i] + 3600.0 * area * np.sum(nc.variables[tracer+'_E'][:])

        with Dataset(ofn) as nc:
            t = nc.variables[tracer][:] # kg/kg

            # convert kg/kg -> kg
            p = nc.variables['P'][:]
            ps = nc.variables['PS'][:]

            # mass of air (kg/m2)
            mair = c2sm.conv.calculate_mair(p, ps, h) * area

            transport[i] = np.sum(t * mair)

        if i > 20:
            break


    return emissions, transport









def int2lm_io(root='.'):
    """\
    Compares input and output emission of INT2LM
    """
    root = '/scratch/snx1600/gkuhl/smartcarb/mother/2015010100_0_72/int2lm'

    inp_filenames = os.path.join(root, 'input', 'emissions', 'emis_*.nc')
    #inp_filenames = os.path.join(root, 'emissions_old', '*.nc')
    print(inp_filenames)
    inp_filenames = sorted(glob(inp_filenames))

    out_filenames = os.path.join(root, 'output', 'lbfd*e.nc')
    #out_filenames = os.path.join(root, 'emissions', '*.nc')
    print(out_filenames)
    out_filenames = sorted(glob(out_filenames))

    print(len(inp_filenames), len(out_filenames))

    message = 'Warning: INT2LM input/output mismatch for "%s" with difference %f!'

    for ifn, ofn in zip(inp_filenames, out_filenames):

        print(os.path.basename(ifn), os.path.basename(ofn))

        with Dataset(ifn) as di, Dataset(ofn) as do:

            keys = [key for key in do.variables.keys() if key[:2] in ['CO', 'NO']]

            for key in keys:
                vi = np.sum(di.variables[key][:,:,2:-2,2:-2])
                #vi = np.sum(di.variables[key][:,:])
                vo = np.sum(do.variables[key][:])

                if abs(vi) < 1e-12:
                    d = abs(vo - vi)
                else:
                    d = abs(vo - vi) / vi

                if d > 1e-4:
                    print(message % (key, d))



if __name__ == '__main__':
    tracer_mass_ts()
    #main("/project/d73/input/")



