#!/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 *

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

    def __init__(self, files, time_step, time_granularity):
        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
            
    """
    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()

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

            # Collect all timestamps
            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            
            for d1 in sorted(dates):                            
              if d0:        
                tok, tdiff = self.get_time_difference(d0, d1)        
                if not(tok):
                  self.logger.error('Temporal inconsistency [t0={} - t1={}] = {} sec'.format(d0, d1, tdiff))  
                res = res & tok                           
              d0 = d1

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

    def main(self):

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

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

