#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)

from datetime import date, datetime, timedelta
from calendar import monthrange
from cdo import Cdo

# Local imports
from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_bokeh import *

"""
Class: Temporal consistency checks
---
- Check frequency/leadtimes of a given variable
"""
class TemporalConsistencyTest:

    def __init__(self, resdir, testname, files, time_step, time_granularity, is_complete):
        self.testid = 'test05-temporal_consistency'
        self.logger = logging.getLogger(self.testid)
        self.time_s = float(time_step)
        self.time_g = time_granularity.lower()                
        self.files = files
        self.complete = (is_complete.strip() == 'yes')
        self.plot_title = testname.replace('.ini', '')
        self.outfname_plot = os.path.join(resdir, testname.replace('.ini', '-temporal_consistency.html'))

    """
    Difference between two dates in months (takes into account calendar)
    """
    def monthdelta(self, d1, d2):
        delta = 0
        while True:
            mdays = monthrange(d1.year, d1.month)[1]
            d1 += timedelta(days=mdays)
            if d1 <= d2:
                delta += 1
            else:
                break
        return delta

    """
    Compute time difference given t1,t2 datetime objects and conditions (time_unit and granularity)
    """
    def get_time_difference(self, t1, t2):

        # Substract datetimes
        tdiff = (t2 - t1).total_seconds()          
        if tdiff == 0:
            return True, 0 # grib messages have multiple timestamps (same date)
        
        # Depending on granularity
        if self.time_g == 'm' or 'month' in self.time_g:
            return self.monthdelta(t1, t2) == self.time_s, tdiff
        elif self.time_g == 'd' or 'day' in self.time_g:
            return (tdiff / 86400) == self.time_s, tdiff
        elif self.time_g == 'h' or 'hour' in self.time_g:
            return (tdiff / 3600) == self.time_s, tdiff
        elif self.time_g == 's' or 'sec' in self.time_g:
            return (tdiff) == self.time_s, tdiff            
        else:
            raise Exception('Please verify the time_granularity parameter (Month, Day, Hour, Sec)')

    def check_timestamps(self):

        # For all files
        cdo = Cdo()
        res = True
        
        # Progress init
        i = 0
        n = len(self.files)
        print()
        errors = ''
        total_dates = []

        for f in self.files:
            # Progress update            
            i+=1         
            print_progress(i, n)            
            
            try:
                info = cdo.showtimestamp(input=f)
            except:
                self.logger.error('CDO-error Bad file format ({})'.format(os.path.basename(f)))
                res = False
                break

            # Collect all timestamps including HHMMSS
            dates = []
            for inf in info:                
                for t in inf.split():
                    match = re.search(r'\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}', t)
                    if match:  # YYYY-MM-DDTHH:MM:SS minimum requirement 
                        dates.append(datetime.strptime(match.group(), '%Y-%m-%dT%H:%M:%S'))

            # Sort and seek
            d0 = False
            dates = sorted(dates)
            total_dates += dates

            for d1 in dates:
                if d0:
                    tok, tdiff = self.get_time_difference(d0, d1)
                    if not(tok):
                        errors = errors + '[{} - {}] = {} sec'.format(d0, d1, tdiff) + ', '
                    res = res & tok
                # Next
                d0 = d1

        # Plot (for visual inspection)
        if len(total_dates) > 1:
            tp = TemporalPlot(self.outfname_plot, self.plot_title)
            tp.plot_consistency(total_dates)

        # Dump errors
        self.logger.error(errors)

        # Final message
        if res:
            msg = 'Files are temporally consistent ({}, {})'.format(self.time_s, self.time_g)
            self.logger.info(msg)
            return 'ok', msg
        else:
            # Monthly datasets with 100% download strategy
            if self.complete:
                msg = 'Temporal inconsistency found. Check log messages in the following pages'
                self.logger.error(msg)
                return 'err', msg

            # Daily datasets with random download strategy (2 months x year)
            else:
                msg = 'Temporal inconsistency found. Reminder: the dataset has not been downloaded completely'
                self.logger.warning(msg)
                return 'warn', msg

    def main(self):

        # Do the actual work
        res, msg = self.check_timestamps()

        # Prepare output
        return { 
            'res': res, 
            'msg': msg
        }        

