#!/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)
#       Ivan Jose Cernicharo Ortiz (ivan.cernicharo@bsc.es)

import sys
import logging
import os
import json
import re
import subprocess
import calendar
from datetime import datetime, date, time, timedelta
from cdo import Cdo

# Local import
from utils.utils_files 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 temporal_completeness_test:

    def __init__(self, logdir, files, forms_dir, dataset, variable, origin, system, is_seasonal):
        self.testid = 'test03-temporal_completeness'
        self.logger = logging.getLogger(self.testid)
        self.logdir = logdir 
        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

    def get_forms_dates(self):
        constraints = {}
        
        # Load constraints file into dictionary
        try:            
            with open(os.path.join(self.forms_dir, self.dataset + '/constraints.json'), 'r') as cdf:
                constraints = json.load(cdf)                
        except:
            raise Exception('Problem loading constraints from cds-forms directory. Please check your configuration')

        # In particular dates for a given (dataset, variable)
        md = []
        try:
            dates = []
            for c in constraints:

                # Filter by variable
                if self.variable in c['variable']:

                    # Seasonal needs filter by origin/system
                    if 'seasonal' in self.dataset:
                        if not(self.origin in c['originating_centre']) or not(self.system in c['system']):
                            continue

                    # Parse all dates
                    for y in c['year']:
                        for m in c['month']:
                            if 'day' in c:
                                for d in c['day']:
                                    dates.append(datetime(int(y), int(m), int(d)))
                            else:
                                dates.append(datetime(int(y), int(m), int(1)))

            # Get unique ordered values
            md = sorted(list(set(dates)))
        except Exception as e:
            print(e)
            raise Exception('Problem loading constraints from cds-forms directory. Please check your dataset name')

        return md



    def check_timestamps(self):

        # Get mandatory dates from cds-forms
        mandatory_dates = self.get_forms_dates()
        dates_available = []

        # For all files
        cdo = Cdo()
        res = 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('Bad file format (cdo cannot read file)')
                print('ERROR: Bad file format ({})'.format(os.path.basename(f)))
                res = False
                break

            # Gather all available timestamps 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        
        for md in mandatory_dates:            
            # Date is not there (date can be there in multiple lead times)
            if dates_available.count(md) == 0:
                res = False
                self.logger.error('Date {} was not found in the provided directory'.format(md))
                
        # Final message
        if res:
            msg = 'Files are temporally complete (according to c3s-forms)'
            self.logger.info(msg)            
            return 'ok', msg
        else:
            msg = 'Temporal incompleteness found. Check log file for more details'
            self.logger.info(msg)            
            return 'err', 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()