from __future__ import print_function
from builtins import str
from builtins import range
from builtins import object
import numpy as np
from datetime import date, datetime, timedelta
from qamods.helpers import conv2jul, parse_ratio
import qamods.io.fauxioapi as io

class data_out(object):
    """
    Outfile class.
    """

    def __init__(self, outfile_name, out_type, gsdate, verbosity = False):
        self.outfile_name = outfile_name
        self.out_type = out_type.upper()
        self.gsdate = gsdate
        self.verbosity = verbosity    
        if self.verbosity == True: 
            print('Opening %s file for writing: %s' %(self.out_type, self.outfile_name))
        if self.out_type == 'NCF':
            self.outfile = self._open_NCF()
        elif self.out_type == 'CSV':
            self.outfile = self._open_CSV()
        else:
            raise ValueError('Wrong outfile type specified.')

    def __str__(self):
        return self.outfile_name

    def __call__(self):
        return self.outfile

    def write_outfile(self, out_dict, grid = '', grid_desc = '', region = '', tons = False, units = '',
       srg_file=''):
        species_list = list(out_dict.keys())
        while '0' in species_list:  # Dunno why this gets inserted?
            species_list.remove('0')
        self.species_list = sorted(species_list)
        if self.out_type == 'NCF':
            self._write_NCF(out_dict, self.species_list, grid, grid_desc, tons, units)
        elif self.out_type == 'CSV':
            if region == 'state' or region == 'county':
                self._write_FIPS_CSV(out_dict, self.species_list, grid, grid_desc, tons, region, srg_file)
            else: 
                self._write_grid_CSV(out_dict, self.species_list, tons)

    def _open_CSV(self):
        try: 
            outfile = open(self.outfile_name, 'w')
        except IOError:
            raise IOError('%s not available for access.' %self.outfile_name)
        else: 
            return outfile 

    def _write_grid_CSV(self, out_dict, species_list, tons):
        '''
        Writes the dictionary to an output file in csv format by grid cell.  Takes output dictionary and the species_list.
        '''
        if out_dict[species_list[0]]().shape[0] == 1:
            hour_type = 'sum'
        else:
            hour_type = 'hourly' 
        self.outfile.write('hour,row,col,' + ','.join(species_name for species_name in species_list) + '\n')
        for hour in range(out_dict[species_list[0]]().shape[0]):
            if hour_type == 'sum':
                out_hour = 'sum'
            else:
                out_hour = hour
            for row in range(out_dict[species_list[0]]().shape[1]):
                srow = str(row)
                for col in range(out_dict[species_list[0]]().shape[2]):
                    scol = str(col) 
                    out_line = '%s,%s,%s' %(out_hour, row + 1, col + 1)
                    for species_name in species_list:
                        out_line = '%s,%s' %(out_line, out_dict[species_name]()[hour,row,col])
                    self.outfile.write(out_line + '\n')

    def _write_FIPS_CSV(self, out_dict, species_list, grid, grid_desc, tons, region, srg_file):
        '''
        Writes the dictionary to an output file in csv format by fips.  Takes output dictionary and the species_list.
        '''
        if not grid: 
            raise ValueError('No grid specified.  Grid needed to write state or county based csv.')
        grid_info = io.Grid(grid, grid_desc)
        ratio_table = parse_ratio(region, grid_info, srg_file)
        self.outfile.write('hour,fips,' + ','.join(species_name for species_name in species_list) + '\n')
        fips_list = sorted(ratio_table.keys())
        if out_dict[species_list[0]]().shape[0] == 1:
            hour_type = 'sum'
        else:
            hour_type = 'hourly' 
        for hour in range(out_dict[species_list[0]]().shape[0]):
            if hour_type == 'sum':
                out_hour = 'sum'
            else:
                out_hour = hour
            for fips in fips_list:
                out_line = '%s,%s' %(out_hour, fips)
                line_dict = dict((species_name,float(0)) for species_name in species_list)
                for species_name in species_list:
                    for cell in ratio_table[fips]:
                        col = int(cell.split(',')[0])
                        row = int(cell.split(',')[1])
                        if row > grid_info.NROWS or col > grid_info.NCOLS:
                            continue  # Skip ratios that are outside the domain
                        line_dict[species_name] = line_dict[species_name] + (out_dict[species_name]()[hour,row - 1,col - 1] * ratio_table[fips][cell])
                    out_line = '%s,%s' %(out_line, line_dict[species_name])
                self.outfile.write(out_line + '\n') 

    def _open_NCF(self):
        '''
        Opens the netCDF input file and returns an open file object.
        '''
        try: 
            outfile = io.IODataset(self.outfile_name, 'w')
        except TypeError:
            raise IOError('%s not available for access.' %self.outfile_name)
        else: 
            return outfile 

    def _write_NCF(self, out_dict, species_list, grid, grid_desc, tons = False, units = ''):
        '''
        Writes the dictionary to an output file in NCF format.  Takes output dictionary and the species_list.
        '''
        hours = out_dict[species_list[0]]().shape[0]
        grid_info = io.Grid(grid, grid_desc)
        self.outfile.set_dimensions('GRID', LAY=1, ROW=grid_info.NROWS, COL=grid_info.NCOLS, VAR=len(species_list))
        for species_name in species_list:
            self._write_species(out_dict, species_name, tons, units)
        self._outfile_settings(hours, grid_info)
        self.outfile.write_TFLAG()
        self.outfile.close()
            
    def _write_species(self, out_dict, species_name, tons = False, units = '', long_name = '', var_desc = ''):
        """
        Takes the output dictionary, species name, and optionally long name, units, and variable description.
        Creates a species of name species_name with the standard smoke shape of TSTEP, LAY, ROW, and COL.
        Returns a multidimensional array of shape TSTEP, LAY, ROW, COL.
        """
        if not units:
            if tons:
                units = 'tons/day'
            else:
                units = 'moles/s'
        if not long_name: 
            long_name = species_name
        if not var_desc: 
            var_desc = 'Model species ' + species_name
        d_shape = [out_dict[species_name]().shape[0],1,out_dict[species_name]().shape[1],out_dict[species_name]().shape[2]]
        species_out = self.outfile.create_variable(species_name, 'REAL', ('TSTEP','LAY','ROW','COL'), 
          long_name=long_name, units=units, var_desc=var_desc)
        data_out = np.zeros(d_shape, np.float32)
        for lay in range(data_out.shape[1]):
                try:
                    data_out[:,lay,:,:] = out_dict[species_name]()
                except ValueError:
                    raise ValueError('Array size mismatch. Please check that input domain matches the size of the intended output domain (-G [GRID]).')
        species_out[:] = data_out

    def _outfile_settings(self, hours, grid):
        '''
        Set the output file dimensions and IOAPI metadata
        '''
        # Set the time step based on how many hours are in the DS
        if hours > 1:
            hstep = 10000
        else:
            hstep = 240000
        esdate = conv2jul(self.gsdate)
        self.outfile.set_attributes(esdate, grid, FILEDESC='EMISQA'.ljust(80), TSTEP=hstep)

