#!/usr/bin/env python

from __future__ import print_function

import sys
import os
import string
import datetime
import yaml
import glob
import re
import argparse
import pandas as pd
import numpy as np
import netCDF4 as nc
import matplotlib.pyplot as plt

"""
Compare the outputs of two or more experiments. 

This script will look at model output, calculate metrics for fields, and
compare/plot these between two runs/experiments. 

It uses the payu experiment layout to find inputs, although it should be easy
to use a different layout. 

FIXME: there's lots that can be done to speed things up.
    - Don't calculate timeseries for points that are not overlapping. 
    - Perhaps have a cache of timeseries that are hashed using the time
      periods. 
"""

comparison_html = string.Template("""
<!DOCTYPE html>
<html>
    <body>
        $plots
    </body>
</html>
""")

class Field:

    def __init__(self, name, label, files, metric,
                 area_of_interest, expected_value=None):
        self.name = name
        self.label = label
        self.metric = metric
        self.files = files
        self.aof = area_of_interest
        self.expected_value = expected_value
        self.ts = None
        self.long_name = None
        self.units = None

        with nc.Dataset(files[0]) as f:
            try:
                self.long_name = f.variables[name].long_name 
                self.units = f.variables[name].units
            except AttributeError:
                pass


    def calc_timeseries(self, time_points=None):
        """
        Calculate a timeseries.
        """

        data = []
        index = []

        for file in self.files:
            with nc.Dataset(file) as f: 
                var = f.variables[self.name]
                assert(var.dimensions[0] == 'time')

                # Get the start date. 
                time_var = f.variables['time']
                try:
                    m = None
                    m = re.search('\d{4}-\d{2}-\d{2}', time_var.units)
                except AttributeError:
                    pass
                if m is not None:
                    start_date = datetime.datetime.strptime(m.group(0), '%Y-%m-%d')
                else:
                    start_date = datetime.date(1, 1, 1)

                for (i, time_units) in enumerate(time_var):

                    if time_points is not None and i < time_points[0]:
                        continue
                    if time_points is not None and i > time_points[1]:
                        break

                    # Select the area of interest. Only handles up to 3
                    # dimensions.
                    if len(var.shape) == 2:
                        val = eval(self.metric)(var[i, self.aof[0]:self.aof[1]])
                    elif len(var.shape) == 3:
                        val = eval(self.metric)(var[i, self.aof[0]:self.aof[1],
                                              self.aof[2]:self.aof[3]])
                    elif len(var.shape) == 4:
                        val = eval(self.metric)(var[i, :, self.aof[0]:self.aof[1],
                                                    self.aof[2]:self.aof[3]])
                    else:
                        assert(False)

                    val = np.atleast_1d(val)
                    assert(len(val) == 1)
                    data.append(val[0])

                    # FIXME: hack for masked time values. Fix dump files. 
                    masked = False
                    try:
                        masked = time_units.mask
                    except AttributeError:
                        pass

                    if masked:
                        index.append(i)
                    else:
                        delta = datetime.timedelta(days=int(time_units))
                        if 'seconds' in time_var.units:
                            delta = datetime.timedelta(seconds=int(time_units))
                        d = start_date + delta
                        d = (d - datetime.datetime(1, 1, 1)).total_seconds()
                        index.append(d)

                assert(len(data) == len(index))

        # Make sure there are no duplicate indices.
        assert(len(set(index)) == len(index))  
        s = pd.Series(data, index)
        self.ts = s.sort_index()


def plot_fields(fields, plot_name, title):
    """
    Simple plot of timeseries of multiple field metrics.
    """

    if any([len(f.ts) == 0 for f in fields]):
        return None

    def get_overlapping_range(intervals):
        """
        Look at a collection of intervals, return the interval on which they
        all overlap.

        The input is a list of tuples [(begin, end), (begin, end) ...]
        """

        def get_overlapping((b0, e0), (b1, e1)):
            """
            Get overlapping of two ranges. 
            """

            begin = max(b0, b1)
            end = min(e0, e1)

            if begin > end:
                return (None, None)
            else:
                return (begin, end)
                

        int_b, int_e = intervals[0]
        for b, e in intervals:
            
            (o_b, o_e) = get_overlapping((int_b, int_e), (b, e))

            if o_b is not None:
                int_b = o_b
                int_e = o_e

        return int_b, int_e

    begin, end = get_overlapping_range([(f.ts.index[0], f.ts.index[-1]) for f in fields])
    if begin is None:
        return None

    plt.figure()
    if title is not None:
        plt.title(title)
    else:
        plt.title(fields[0].long_name)
    plt.xlabel('time (seconds)')
    plt.ylabel(fields[0].units)
    for f in fields:
        f.ts[begin:end].plot(label=f.label)
    plt.legend()

    if plot_name is None:
        plot_name = (fields[0].name).replace(' ', '_')
        
    plot_name = plot_name + '.png'
    plt.savefig(plot_name)

    return plot_name


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('comparison_files', nargs='+',
                        help='A list of files defining the comparisons to be made.')
    parser.add_argument('--output_dir', default='./',
                        help='Directory where plots and data will go.')

    args = parser.parse_args()
    page_name = 'field_comparison.html'

    img_str = ''
    for file in args.comparison_files:
        with open(file, 'r') as f:
            doc = yaml.load(f)
            page_name = doc['page_name']

            for comp in doc['comparisons']:
                fields = []
                for f in comp['fields']:
                    files = glob.glob(f['file'])
                    if files == []:
                        sys.stderr.write('Not found: {} \n'.format(f['file']))
                        return 1

                    aof = f.get('area_of_interest',
                                [0, sys.maxint, 0, sys.maxint])
                    metric = comp.get('metric', 'lambda x : x')

                    field = Field(f['name'], f['label'], files, metric, aof)
                    time_points = comp.get('time_points')
                    if time_points is not None:
                        if type(time_points) != type([]):
                            time_points = [0, time_points]
                    field.calc_timeseries(time_points)
                    fields.append(field)

                plot_name = plot_fields(fields, comp.get('plot_name'),
                                        comp['title'])
                if plot_name is not None:
                    img_str += "<img src=\"{}\">\n".format(plot_name)

    # Now make a new html page with the new plots. 
    page = comparison_html.substitute(plots=img_str)
    with open(page_name, 'w') as f:
        f.write(page)


if __name__ == '__main__':
    sys.exit(main())

