
from datetime import datetime
import calendar
import itertools
import os
import random
import string

from matplotlib.dates import DateFormatter
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import scipy.stats
import pandas

import c2sm

from smartcarb import misc
from smartcarb import map
from smartcarb import config as cfg


import dplume.io


DLON_ORBIT = 24.146341463414636

COLORS = [
    '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2',
    '#7f7f7f', '#bcbd22', '#17becf'
]


def plot_signal_and_background(origin='Berlin'):
    """
    Plot CO2, NO2 and CO signals and background (mean and variance) for
    Berlin.
    """
    if origin == 'Berlin':
        suffix = 'Berlin_100'
    else:
        suffix = '%s_50' % origin

    path = '/project/smartcarb/data1/ts/'
    bgCO2 = pandas.read_csv(path+'ts_XCO2_background_%s.dat' % suffix,
                            delimiter='\t', index_col=0, parse_dates=True)
    bgCO2 = bgCO2[bgCO2.index.hour == 11]

    bgNO2 = pandas.read_csv(path+'ts_NO2_background_%s.dat' % suffix,
                            delimiter='\t', index_col=0, parse_dates=True)
    bgNO2 = bgNO2[bgNO2.index.hour == 11]

    bgCO = pandas.read_csv(path+'ts_CO_background_%s.dat' % suffix,
                            delimiter='\t', index_col=0, parse_dates=True)
    bgCO = bgCO[bgCO.index.hour == 11]

    path = '/project/kug/smartcarb/results/%s' % origin
    if origin == 'Berlin':
        filename = os.path.join(path, 'Berlin_CO2+NO2+CO_signals_ns37.dat')
    else:
        filename = os.path.join(path, '%s_CO2+NO2+CO_signals_ns5.dat' % origin)

    signals = pandas.read_csv(filename, delimiter='\t', index_col=0,
                              parse_dates=True)
    signals = misc.filter_dataset(signals, nsat=6)

    data = {
        'CO2': [
            signals['CO2_signal'],
            bgCO2['mean_bg'],
            bgCO2['std_bg']
        ],
        'NO2': [
            signals['NO2_signal'] / 1e16,
            bgNO2['mean_bg'] / 1e16,
            bgNO2['std_bg'] / 1e16
        ],
        'CO': [
            signals['CO_signal'] / 1e17,
            bgCO['mean_bg'] / 1e17,
            bgCO['std_bg'] / 1e17
        ],
    }
    uncer = {
        'CO2': [0.5, 0.7, 1.0],
        'NO2': np.array([1e15, 2e15]) / 1e16,
        'CO':  [4e17 / 1e17],
    }

    name2ylim = {
        'Berlin': {
            'CO2': [(0, 3.0), (390, 410), (0, 3)],
            'NO2': [(0, 4.0), (0, 4), (0, 4.0)],
            'CO':  [(0, 5.0), (10, 25), (0, 5.0)]
        },
        'Janschwalde': {
            'CO2': [(0, 13.0), (390, 410), (0, 3)],
            'NO2': [(0, 18.0), (0, 4), (0, 4.0)],
            'CO':  [(0, 7.0), (10, 25), (0, 5.0)]
        },
        'Lippendorf': {
            'CO2': [(0, 13.0), (390, 410), (0, 3)],
            'NO2': [(0, 18.0), (0, 4), (0, 4.0)],
            'CO':  [(0, 7.0), (10, 25), (0, 5.0)]
        },
    }[origin]

    name2ylabel = {
        'CO2': 'XCO$_2$ (ppm)',
        'NO2': 'NO$_2$ (10$^{16}$ molec. cm$^{-2}$)',
        'CO':  'CO (10$^{17}$ molec. cm$^{-2}$)'
    }

    tick_locations = {
        'Berlin':      [1,5,1,1,1,1,1,5,1],
        'Janschwalde': [2,5,1,2,1,1,1,5,1],
        'Lippendorf':  [2,5,1,2,1,1,1,5,1],
    }[origin]


    with plt.rc_context({
            'axes.autolimit_mode': 'round_numbers',
            'font.size': 9,
            'xtick.labelsize': 'x-small',
            'ytick.labelsize': 'x-small',
            #'text.usetex': True,
            #'text.latex.unicode': True,
        }):
        fig = plt.figure(figsize=(7,5.5))

        for i, tracer in enumerate(['CO2', 'NO2', 'CO']):
            for j in range(3):

                labels = ['signal', 'total', 'total']
                headings = ['signal', 'background mean',
                            'background standard deviation']

                k = i*3 + j
                ax = fig.add_subplot(3,3,k+1)

                ax.plot(data[tracer][j], label=labels[j])

                # reference noise
                if j == 0:
                    ax.hlines(uncer[tracer], '2015-01-01', '2015-12-31',
                              label='uncer.', colors='#ff7f0e', zorder=3)

                    print(tracer, end='\t')
                    for p in [5,50,95]:
                        s = scipy.stats.scoreatpercentile(data[tracer][j], p)
                        print('%.2f' % s, end='\t')
                    print()

                if j == 2:
                    #print(tracer, end='\t')
                    for p in [5,50,95]:
                        s = scipy.stats.scoreatpercentile(data[tracer][j], p)
                        #print('%.2f' % s, end='\t')
                    #print()

                # biospheric signal
                if tracer == 'CO2' and j == 2:
                    try:
                        ax.plot(bgCO2['std_bio'], label='only (local) biosphere')
                    except KeyError:
                        print('Warning: No biospheric XCO2 ')

                ax.legend(loc=1, fontsize='small')
                ax.grid(True, ls='-', color='silver')

                ax.set_xlim('2015-01-01', '2015-12-31')
                ax.set_ylim(*name2ylim[tracer][j])

                t = tick_locations[k]
                ax.yaxis.set_major_locator(ticker.MultipleLocator(t))

                ax.set_title('(%s) %s' % (string.ascii_letters[k],
                                          headings[j]), fontsize='small')
                #date_labels(ax, fmt='%b', axis='x')

                locs = [datetime(2015,m,1) for m in range(1,13)]
                labels = [calendar.month_abbr[m][0] for m in range(1,13)]
                plt.xticks(locs, labels)

                if j == 0:
                    label = name2ylabel[tracer]
                    ax.text(-.15,0.5, label, transform=ax.transAxes, rotation=90,
                            fontsize='small', va='center')

                    #ax.set_ylabel(name2ylabel[tracer], fontsize='small',
                    #              position=(0,ymin+0.5*(ymax-ymin)),)
                if i == 2:
                    ax.set_xlabel('2015', fontsize='small')

        plt.tight_layout(rect=[0.01,-0.01,1.01,1.01])



def read_detections_mc(filename, column, nsat=6, nsamples=20):

    d = pandas.read_csv(filename, sep='\t', index_col=0, parse_dates=True)
    d = misc.filter_dataset(d, nsat)

    time = np.array(d.index).reshape(-1,nsamples)

    for t in time:
        if np.any(t[0] != t):
            raise ValueError('times of samples are not equal')

    return time[:,0], d[column].values.reshape(-1,nsamples)


def visualize(data, do_zoom=True, show_clouds=True, use_noise=True,
              is_sentinel_5=False, data7=None):

    domain = cfg.DOMAINS['smartcarb']
    lon_o, lat_o, radius = data.attrs['plume location (lon, lat, radius)']

    # tracer field
    if not use_noise:
        noise = 0.0

    clouds = data.is_cloudy if show_clouds else None

    if data.tracer == 'CO2':
        mean_bg = np.nanmean(data.CO2)
        fig = c2sm.vis.make_level2_map(
            data.lonc, data.latc, data.CO2 - mean_bg, vmin=-2, vmax=2,
            label='XCO$_2$ signal (ppm)', clct=clouds, domain=domain
        )
    else:
        vmax = 10 if is_sentinel_5 else 20
        fig = c2sm.vis.make_level2_map(
            data.lonc, data.latc, data.NO2 / 1e15, vmin=0, vmax=vmax, truncate_cmap=True,
            label='NO$_2$ column (10$^{15}$ molecules per cm$^2$)', clct=clouds,
            domain=domain
        )
    #
    ax, cax = fig.get_axes()
    map.add_city_labels(ax)

    # lines and labels
    lines = ax.plot(cfg.pollon, cfg.pollat, marker='s', mec='k', mfc='w', ls='',
            transform=ccrs.PlateCarree())
    labels = ['cloud fraction > %d%%' % (100*clouds.threshold)]

    # true plumes
    for tracer, ls in [('CO2', '-'), ('NO2', '--')]:
        name = '%s_true_plume' % tracer

        if tracer == 'CO2':
            thr = data[name].attrs['signal threshold (ppm)']
            label = 'CO$_2$ plume (>%.2f ppm)' % thr
        elif tracer == 'NO2':
            thr = data[name].attrs['signal threshold (molecules cm-2)']
            label = r'NO$_2$ plume (>%.1f$\times$10$^{15}$ cm$^{-2}$)' % (thr / 1e15)
        else:
            label = name

        if hasattr(data, name):
            lines += c2sm.vis.draw_contours(ax, data.lon, data.lat,
                                            data[name], label=label, ls=ls)
            labels += [label]
            continue

        if hasattr(data7, name):
            lines += c2sm.vis.draw_contours(ax, data7.lon, data7.lat,
                                            data7[name], label=label, ls=ls)
            labels += ['%s (Sentinel-7)' % label]

    # detected plumes
    if data.detected_plume is not None and np.sum(data.detected_plume) > 0:

        lon0 = data.lon.values[data.detected_plume]
        lat0 = data.lat.values[data.detected_plume]

        lines += ax.plot(lon0, lat0, marker='o',
                mec='none', ms=1.5, color='k', ls='',
                transform=ccrs.PlateCarree()
        )
        if hasattr(data, 'confusion_matrix_30'):
            extra = ' (PPV = %.2f)' % data.confusion_matrix_30.PPV
        else:
            extra = ''

        q = data.attrs['probability for z-value']
        ns = data.attrs['size of neighborhood']

        labels += [
            'detected plume (q = %.3f, n = %d)\n' % (q, ns) +
            'with %d pixels%s' % (np.sum(data.detected_plume), extra)
        ]

    if do_zoom:
        try:
            rlon_0, rlat_0 = c2sm.transform.rotpole2wgs(
                data.CO2_true_plume.center_lon,
                data.CO2_true_plume.center_lat,
                cfg.pollon, cfg.pollat, inverse=True)
        except AttributeError:
            rlon_0, rlat_0 = c2sm.transform.rotpole2wgs(lon_o, lat_o, cfg.pollon,
                                              cfg.pollat, inverse=True)

        dy = 1.1
        dx = dy * 1.12
        xmin, xmax = rlon_0 - dx, rlon_0 + dx
        ymin, ymax = rlat_0 - dy, rlat_0 + dy

        try:
            ax.set_xlim(xmin, xmax)
            ax.set_ylim(ymin, ymax)
        except ValueError:
            pass

    if lines:
        ax.legend(lines, labels, loc="upper left", numpoints=1, markerscale=1.5,
                fontsize='small')


    return fig




def detection_limits(filename):
    d = pandas.read_csv(filename, index_col=0)

    fig = plt.figure()
    ax = fig.add_subplot(111)


    valids = np.all(np.isfinite(d.values), axis=1)
    limits = d.values[valids]

    ax.bar(np.arange(5), np.mean(limits, 0), yerr=np.std(limits, 0))

    ax.set_ylim(0,0.6)
    ax.set_ylabel('detection limit (ppm)')
    ax.grid(axis='y')

    ax.set_xticks(np.arange(5))
    ax.set_xticklabels([c.replace('2_', '$_2$\n') for c in d.columns])

    return ax



def month_labels(ax=None):
    if ax is None:
        ax = plt.gca()

    numbers = np.arange(1,13)
    labels = [calendar.month_abbr[n] for n in numbers]
    ax.set_xticks(numbers)
    ax.set_xticklabels(labels)



def date_labels(ax=None, fmt='%b', axis='x'):

    fmt = DateFormatter(fmt)

    if 'x' in axis:
        ax.xaxis.set_major_formatter(fmt)

    if 'y' in axis:
        ax.yaxis.set_major_formatter(fmt)


def plot_obs_per_month_and_scenario(nmax=99, nsat=3, cc=1, show=True):

    products = [
        ('CO2', 0.5, 37),
        ('CO2', 0.7, 37),
        ('CO2', 1.0, 37),
        ('NO2', 'S5', 1),
        ('NO2', 'low', 37),
        ('NO2', 'high', 37),
    ]

    product2label = {
        ('CO2', 1.0): 'CO$_2$, 1.0 ppm',
        ('CO2', 0.7): 'CO$_2$, 0.7 ppm',
        ('CO2', 0.5): 'CO$_2$, 0.5 ppm',
        ('NO2', 'S5'):   'NO$_2$, Sentinel-5',
        ('NO2', 'low'):  r'NO$_2$, 1$\times$10$^{15}$ cm$^{-2}$',
        ('NO2', 'high'): r'NO$_2$, 2$\times$10$^{15}$ cm$^{-2}$'
    }

    if show:
        fig = plt.figure(figsize=[11, 4])
        ax = fig.add_subplot(111)

    lon_eqs = np.linspace(0, DLON_ORBIT, nsat+1)[:-1]

    values = []
    errors = []

    for i,(tracer, ssp, ns) in enumerate(products):

        # estimate uncertainty
        data = dplume.io.read_results(tracer, ssp, ns, nsat=None)
        sample = [
            sum(data['TP%d' % cc][(data.PPV30 >= 0.8) & (data.lon_eq == l)] >= 100)
            for l in sorted(set(data.lon_eq))
        ]
        error = np.std([sum(random.sample(sample, nsat)) for _ in range(5000)])
        errors.append(error)

        data = dplume.io.read_results(tracer, ssp, ns, nsat=nsat)
        months, days = days_per_month(data, lon_eqs, nmax, cc=cc)
        cdays = days.sum(0)

        if show:
            ax.bar(months+i/7.0-0.5+1/7, cdays, width=1/7,
                   label='%s (%d)' % (product2label[tracer,ssp], sum(cdays)))

        values.append(sum(cdays))



    if show:
        month_labels(ax)

        ax.set_xlabel('months in 2015')
        ax.set_ylabel('number of plume observations')

        plt.legend(ncol=3, framealpha=0.4, title='Instrument scenario (number of plumes)')
        plt.tight_layout()
        plt.grid(axis='y')
        ax.set_ylim(0,9)

    # add black line showing the max. number of detectable plumes
    data = dplume.io.read_overpasses('Berlin', nsat=nsat)
    months, days = days_per_month(data, lon_eqs, nmax, cc=cc)
    cdays = days.sum(0)

    values.append(sum(cdays))

    data = dplume.io.read_overpasses('Berlin', nsat=None)
    sample = [
        sum(data['pixels (<1%)'][data.lon_eq == l] >= 100)
        for l in sorted(set(data.lon_eq))
    ]
    error = np.std([sum(random.sample(sample, nsat)) for _ in range(5000)])
    errors.append(error)

    if show:
        for i,month in enumerate(months):
            ax.plot([month-0.43, month+0.43], [cdays[i], cdays[i]], 'k-')


    return values,errors




def plot_plumes_per_sat(days=None):
    sats = np.arange(1,7)

    labels = [
        r'CO$_2$, 0.5 ppm',
        r'CO$_2$, 0.7 ppm',
        r'CO$_2$, 1.0 ppm',
        r'NO$_2$, Sentinel-5',
        r'NO$_2$, 1$\times$10$^{15}$ cm$^{-2}$',
        r'NO$_2$, 2$\times$10$^{15}$ cm$^{-2}$',
        'real plumes'
    ]

    if days is None:
        days = np.array([
            plot_obs_per_month_and_scenario(nmax=99, nsat=i, cc=1, show=False)
            for i in sats
        ])

    fig = plt.figure()
    ax = fig.add_subplot(111)

    for i in range(7):
        if i == 6:
            color = 'k'
        else:
            color = None

        line = ax.errorbar(sats, days[:,0,i], days[:,1,i], marker='o',
                           color=color, ls='', label=labels[i])[0]

        res = scipy.stats.linregress(sats, days[:,0,i])
        print(i, res)

        x = np.linspace(0.8,6.2)
        ax.plot(x, res.slope * x + res.intercept, '-', color=line.get_color())

    ax.set_xlim(0.6, 6.4)
    ax.set_ylim(0,70)

    ax.set_xlabel('number of satellites')
    ax.set_ylabel('number of plumes per year')

    plt.legend()




def plot_plumes_with_more_than_n_pixels(filename):
    # use Berlin_CO2_Plume.dat in results
    d = pandas.read_csv(filename, sep='\t', parse_dates=True, index_col=0,
                        comment='#')

    n = np.arange(0,501,25) + 1

    fig = plt.figure()
    ax = fig.add_subplot(111)

    marker = 2*'osv^<>dDx*'
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c']

    lon_eqs = sorted(set(d.lon_eq))
    for i,nsat in enumerate([1,3,6]):

        pixels = []
        for lon_eq in itertools.combinations(lon_eqs, nsat):
            px = misc.filter_dataset(d, nsat=1, lon_eqs=lon_eq)
            px = px['pixels (<1%)'].values
            px = [sum(px > nn) for nn in n]
            pixels.append(px)

        px_mean = np.mean(pixels, axis=0)
        px_std = np.std(pixels, axis=0)
        px_min = np.min(pixels, axis=0)
        px_max = np.max(pixels, axis=0)

        ax.plot(n, px_mean, '-', label=str(nsat),
                marker=marker[i], mfc=colors[i], mec='black')

        ax.fill_between(n, px_mean - px_std, px_mean + px_std, color='silver')


        print(nsat)
        for nn in [1,50,100,200]:
            j = np.argmax((nn+1) == n)
            print(nn, px_mean[j], px_std[j], px_min[j], px_max[j])
        print()

    ax.set_xlabel('minimum number of pixels')
    ax.set_ylabel('number of plumes')

    ax.set_xlim(0,500)
    ax.set_ylim(0,120)

    ax.grid(True)

    plt.legend(title='number of satellites')
    plt.tight_layout()



def plot_obs_per_month_and_nsat(filename, nmax=99, index_col=0,
                                only_overlaps=False, cloud_thr=0.01,
                                bootstrap=False):
    """
    Plots for different number of satellites, the number of plume observations
    per month with more than nmax XCO2 pixels (cloud fraction < 1%).
    """
    data = pandas.read_csv(filename, index_col=index_col, sep='\t',
                           parse_dates=True, comment='#')

    with plt.rc_context({
            'font.size': 9,
            'xtick.labelsize': 'medium',
            'ytick.labelsize': 'medium',
            #'text.usetex': True,
            #'text.latex.unicode': True,
        }):
        fig = plt.figure(figsize=[11, 3])
        ax = fig.add_subplot(111)

        for i,nsat in enumerate(range(1,7)):

            months, days = days_per_month(data, nsat, nmax, cloud_thr=cloud_thr,
                                          only_overlaps=only_overlaps,
                                          bootstrap=False)


            if bootstrap:
                _, mean, std, attrs = days_per_month(data, nsat, nmax, cloud_thr=cloud_thr,
                                             only_overlaps=only_overlaps,
                                              bootstrap=True)
                #ax.errorbar(months, days, std, marker='o', color='k', ls='', ms=3)

            # add errorbars
            months = months + i / 7.0 - 0.5 + 1.0 / 7.0
            ax.bar(months, days, width=1/7, yerr=std,
                   label='%d (%d$\pm$%d)' % (nsat, sum(days), attrs['std']))

            print(nsat, sum(days), round(attrs['mean']), round(attrs['std']))

        month_labels(ax)

        ax.set_xlabel('months in 2015')
        ax.set_ylabel('number of plume observations')
        ax.set_ylim(0,11)
        ax.set_yticks(np.arange(ax.get_ylim()[1]+1))

        plt.legend(ncol=3, title='number of satellites (number of plumes)')
        plt.tight_layout(pad=0)
        plt.grid(axis='y', ls='-')



def extract_pixels(data, cloud_thr=0.01, only_overlaps=False):

    if 'pixels (<1%)' in data.columns:
        pixels = data['pixels (<%d%%)' % (100 * cloud_thr)]

        if only_overlaps:
            pixels = pixels[data.does_overlap]
    else:
        pixels = data['TP%d' % (100 * cloud_thr)][data.PPV30 >= 0.80]

    return pixels


def days_per_month(data, nsat, nmax=99, scaling=None, cloud_thr=0.01,
                   only_overlaps=False, bootstrap=False):

    if scaling is None:
        scaling = np.ones(13)

    if not bootstrap:
        data = misc.filter_dataset(data, nsat)
        pixels = extract_pixels(data, cloud_thr, only_overlaps)

        days = np.array([
            scaling[m] * sum(pixels[pixels.index.month == m] > nmax)
            for m in range(1,13)
        ])

        return np.arange(1,13), np.array(days)

    else:
        lon_eqs = sorted(set(data.lon_eq))
        cdays = []
        for lon_eq in itertools.combinations(lon_eqs, nsat):
            data2 = misc.filter_dataset(data, nsat=nsat, lon_eqs=lon_eq)
            pixels = extract_pixels(data2, cloud_thr, only_overlaps)

            days = np.array([
                scaling[m] * sum(pixels[pixels.index.month == m] > nmax)
                for m in range(1,13)
            ])
            cdays.append(days)

        attrs = {
            'nsat': nsat,
            'mean': np.sum(cdays, axis=1).mean(),
            'std':  np.sum(cdays, axis=1).std()
        }
        return np.arange(1,13), np.mean(cdays, axis=0), np.std(cdays, axis=0), attrs




def plot_obs_per_month_boxplot(filename, nsat=4, nmax=-1, use_percent=False,
                               add_sub_constellations=False):

    # time, lon_eq, cloud-free pixels
    #data = pandas.read_csv(filename, index_col=0, parse_dates=True)
    data = filename

    lon_eqs = np.linspace(0, DLON_ORBIT, nsat+1)[:-1]

    if add_sub_constellations and nsat in [4,6]:
        lon_eqs = np.concatenate([lon_eqs[::2], lon_eqs[1::2]])

    if use_percent:
        mdays = np.array(calendar.mdays)
        mdays[12] = 25.0 # only 25 days have been simulated in Dec.
        scaling = 100.0 / mdays
    else:
        scaling = np.ones(13)

    months, days = days_per_month(data, nsat, nmax)

    fig = plt.figure()
    ax = fig.add_subplot(111)

    for i,lon_eq in enumerate(lon_eqs):
        label = '%s (%d°)' % ('ABCDEF'[i], round(lon_eq))
        bottom = np.sum(days[:i], axis=0)
        ax.bar(months, days[i], bottom=bottom, label=label)



    if add_sub_constellations and nsat in [4,6]:
        ax.set_ylim(0,10)
        for i in [1, 3 if nsat == 6 else 2, nsat]:
            pass
            #cdays = days[:i].sum(0)
            #ax.plot(months, cdays - 0.12, 'kv')4
            #for j,m in enumerate(months):
            #    ax.plot([m-0.4, m+0.4], [cdays[j], cdays[j]], 'k-', lw=2)



    month_labels(ax)

    ax.set_xlabel('months in 2015')
    ax.set_ylabel('number of plume observations')

    if add_sub_constellations and nsat in [4,6]:
        plt.legend(ncol=2, title='satellite')
    else:
        plt.legend(ncol=nsat, title='satellites')

    plt.tight_layout()

    return data












