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

import sys
import calendar
from datetime import timedelta

from cdo import Cdo

# Local import
from dqc_wrapper.utils.utils_forms import *
from dqc_wrapper.utils.utils_bokeh import *

"""
Add a month to a give datetime object
"""
def add_month(sourcedate, months=1):
    month = sourcedate.month - 1 + months
    year = sourcedate.year + month // 12
    month = month % 12 + 1
    day = min(sourcedate.day, calendar.monthrange(year,month)[1])
    return datetime(year, month, day)

"""
Add a day to a give datetime object
"""
def add_day(sourcedate, days=1):
    return sourcedate + timedelta(days=1)

"""
Class: Temporal completeness checks
---
- Uses CDO python bindings
- Loops over all timestamps within the dataset
- Checks they are within the time range [start/end] defined as input
"""
class TemporalCompletenessTest:

    def __init__(self, resdir, testname, files, forms_dir, dataset, variable, origin, system, is_seasonal, is_complete):
        self.testid = 'test03-temporal_completeness'
        self.logger = logging.getLogger(self.testid) 
        self.files = files
        self.forms_dir = forms_dir 
        self.dataset = dataset
        self.variable = variable        
        self.origin = origin
        self.system = system
        self.is_seasonal = is_seasonal
        self.complete = (is_complete.strip() == 'yes')
        self.plot_title = testname.replace('.ini', '')
        self.outfname_plot = os.path.join(resdir, testname.replace('.ini', '-temporal_completeness.html'))

    """
    Check timestamps vs mandatory dates from the CDS forms repository
    """
    def check_timestamps(self):

        # Get mandatory dates from cds-forms
        mandatory_dates = get_forms_dates(self.dataset, self.variable, self.origin, self.system, self.forms_dir)
        dates_available = []

        # For all files
        cdo = Cdo()
        ok = True
        
        # Progress init
        i = 0
        n = len(self.files)
        print()

        # Scan for available 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)))
                ok = False
                break

            # Gather all available dates from file
            dates = []
            for inf in info:                
                for t in inf.split(): 
                    match = re.search(r'\d{4}-\d{2}-\d{2}', t)
                    if match:  # YYYYMMDD minimum requirement                                              
                        dates.append(datetime.strptime(match.group(), '%Y-%m-%d'))

            # if seasonal, only the first one to avoid lead_times duplication
            if self.is_seasonal == 'yes':
                dates_available.append(dates[0])
            else:
                for d in dates:
                    dates_available.append(d)

        # First date on seasonal products
        if self.is_seasonal == 'monthly':
            mandatory_dates = [ add_month(m) for m in mandatory_dates ]
        if self.is_seasonal == 'daily':
            mandatory_dates = [ add_day(m) for m in mandatory_dates ]

        # Check consistency
        errors = ''
        missd = 0
        for md in mandatory_dates:            
            # Date is not there (date can be there in multiple lead times)
            if dates_available.count(md) == 0:
                ok = False
                errors = errors + str(md) + ', '
                missd += 1

        # Dump counts
        self.logger.error('Dates missing: {}'.format(missd))

        # Plot (for visual inspection)
        if len(dates_available) > 1 and len(mandatory_dates) > 1:
            tp = TemporalPlot(self.outfname_plot, self.plot_title)
            tp.plot_completeness(dates_available, mandatory_dates)

        # Final message
        if ok:
            msg = 'Files are temporally complete (according to c3s-cds-forms repository)'
            self.logger.info(msg)
            res = 'ok'
        else:
            if self.complete:
                msg = 'Temporal incompleteness found. Check log messages in the following pages'
                self.logger.error(msg)
                res = 'err'
            else:
                msg = 'Temporal incompleteness found. Dataset has not been downloaded completely'
                self.logger.warning(msg)
                res = 'warn'

        # Dump errors
        self.logger.error(errors)
        return res, msg

    def main(self):

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

        # Prepare output
        return { 
            'res': res, 
            'msg': msg
        }
                
if __name__ == "__main__":
    sys.exit()