#!/usr/bin/python3
"""Figures of pikobs.vdedr.

Three figures per selection, with the radiance channels on the y axis:

* ``omp_bias``  residual bias and raw bias, experience minus control
* ``omp_rel``   change of sigma(OMP) and of the number of observations, %
* ``oma_rel``   the same for OMA

Everything is computed from the sums stored by the extraction, so the
SQLite STDDEV extension is not needed any more and the numbers recombine
exactly over cycles.
"""

import os
import re
import sqlite3
import traceback
from typing import Any, Dict, List, Optional, Sequence, Tuple

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

# metric -> (file prefix, departure shown, legend)
METRICS = {
    'omp_bias': ('OMP', ('residual bias', 'raw bias')),
    'omp_rel':  ('OMP', ('% sigma', '% Nobs')),
    'oma_rel':  ('OMA', ('% sigma', '% Nobs')),
}

ROUGE = '#FF9999'
ROUGEPUR = '#FF0000'
VERT = '#009900'
BLEU = '#1569C7'
NOIR = '#000000'
COULEURS = [BLEU, ROUGEPUR, ROUGE, VERT, NOIR]

# Sign of the change, same reading as in scatter: the experience
# reduces the spread (negative) in red, the control is the better one
# (positive) in blue.
EXP_BETTER = "#B2182B"     # red
CTL_BETTER = "#2166AC"     # blue
NEUTRAL_LINE = "#9e9e9e"
BAND = "#f2f2f2"

# Layout of the channel axis. A radiance family can carry more than a
# thousand channels and every one of them must stay readable, so the
# figure grows with the number of channels instead of squeezing them.
INCH_PER_CHANNEL = 0.13
MIN_HEIGHT = 8.0
MAX_HEIGHT = 260.0
TICK_FONTSIZE = 6
DPI = 90
# Observation counts are written next to each channel; past this many
# channels they turn into noise and are dropped.
MAX_CHANNELS_WITH_COUNTS = 300
# Past this many channels the confidence figures are dropped and only the
# significance squares are drawn.
MAX_CHANNELS_WITH_CONFIDENCE = 300

# F-test on the spread, same recipe as saska40_v0_hs.py: the larger
# variance on top, two-sided p, confidence (1 - p) x 100, significant
# above this value.
MIN_CONFIDENCE = 95.0

GRAPHE_NOMVAR = {
    11215: 'U COMPONENT OF WIND (10M)',
    11216: 'V COMPONENT OF WIND (10M)',
    12004: 'DRY BULB TEMPERATURE AT 2M',
    10051: 'PRESSURE REDUCED TO MEAN SEA LEVEL',
    10004: 'PRESSURE',
    12203: 'DEW POINT DEPRESSION (2M)',
    12001: 'TEMPERATURE/DRY BULB',
    11003: 'U COMPONENT OF WIND',
    11004: 'V COMPONENT OF WIND',
    12192: 'DEW POINT DEPRESSION',
    12163: 'BRIGHTNESS TEMPERATURE',
    15036: 'ATMOSPHERIC REFRACTIVITY',
    11001: 'WIND DIRECTION',
    11002: 'WIND SPEED',
    11011: 'WIND DIRECTION AT 10M',
    11012: 'WIND SPEED AT 10M',
}


def varno_name(varno) -> str:
    try:
        varno = int(varno)
    except (TypeError, ValueError):
        return str(varno)
    if varno in GRAPHE_NOMVAR:
        return GRAPHE_NOMVAR[varno]
    try:
        import pikobs
        name, units, _ = pikobs.type_varno(varno)
        return f"{name} {units}".strip()
    except Exception:
        return str(varno)


def _safe(text) -> str:
    return re.sub(r'[^A-Za-z0-9._+-]', '_', str(text))


def ftest_confidence(s_ref, n_ref, s_exp, n_exp) -> np.ndarray:
    """Two-sided F-test on two standard deviations, confidence in percent.

    Same recipe as saska40_v0_hs.py: the larger variance goes on top,
    F = s_big**2 / s_small**2 with (n_big - 1, n_small - 1) degrees of
    freedom, p = min(1, 2 * P(F' > F)), confidence = (1 - p) * 100. The
    standard deviations are the sample ones (n - 1).
    """
    s_ref, s_exp = np.asarray(s_ref, float), np.asarray(s_exp, float)
    n_ref, n_exp = np.asarray(n_ref, float), np.asarray(n_exp, float)
    conf = np.full(s_ref.shape, np.nan)
    ok = (n_ref > 1) & (n_exp > 1) & (s_ref > 0) & (s_exp > 0)
    if not ok.any():
        return conf
    exp_big = s_exp[ok] >= s_ref[ok]
    s_big = np.where(exp_big, s_exp[ok], s_ref[ok])
    s_small = np.where(exp_big, s_ref[ok], s_exp[ok])
    d_num = np.where(exp_big, n_exp[ok], n_ref[ok]) - 1.0
    d_den = np.where(exp_big, n_ref[ok], n_exp[ok]) - 1.0
    F = s_big ** 2 / s_small ** 2
    try:
        from scipy.stats import f as f_dist
        upper = f_dist.sf(F, d_num, d_den)
    except ImportError:
        from math import erfc, sqrt
        z = np.log(F) / np.sqrt(2.0 / d_num + 2.0 / d_den)
        upper = 0.5 * np.vectorize(lambda x: erfc(x / sqrt(2.0)))(z)
    conf[ok] = (1.0 - np.minimum(1.0, 2.0 * upper)) * 100.0
    return conf


def _sample_std(std_pop, n):
    """Population standard deviation (divided by n) -> sample one (n - 1)."""
    n = np.asarray(n, float)
    with np.errstate(divide='ignore', invalid='ignore'):
        return np.where(n > 1, np.asarray(std_pop, float)
                        * np.sqrt(n / np.maximum(n - 1.0, 1.0)), np.nan)


def _channel_stats(db_file: str, region: str, flag: str, stn_sql: str,
                   varno: int) -> Dict[float, Dict[str, float]]:
    """Mean, sigma, count and mean bias correction of every channel.

    stn_sql is the station condition of the selector (empty for join), so
    several platforms are pooled by summing their moments.
    """
    conn = sqlite3.connect(db_file)
    try:
        rows = conn.execute(
            f"""
            SELECT vcoord, SUM(Ntot), SUM(s_omp), SUM(s2_omp),
                   SUM(s_oma), SUM(s2_oma), SUM(n_bcorr), SUM(s_bcorr)
            FROM serie_vdedr
            WHERE region = ? AND flag = ? AND varno = ?{stn_sql}
            GROUP BY vcoord;
            """, (region, flag, int(varno))).fetchall()
    except sqlite3.Error:
        return {}
    finally:
        conn.close()

    out: Dict[float, Dict[str, float]] = {}
    for vcoord, n, s_omp, q_omp, s_oma, q_oma, n_bc, s_bc in rows:
        if not n:
            continue
        n = float(n)
        stat = {'n': n}
        for tag, s, q in (('omp', s_omp, q_omp), ('oma', s_oma, q_oma)):
            if s is None:
                stat[f'avg_{tag}'] = np.nan
                stat[f'std_{tag}'] = np.nan
                continue
            mean = s / n
            var = max((q or 0.0) / n - mean * mean, 0.0)
            stat[f'avg_{tag}'] = mean
            stat[f'std_{tag}'] = np.sqrt(var)
        stat['bcorr'] = (s_bc / n_bc) if (n_bc and s_bc is not None) else np.nan
        out[float(vcoord)] = stat
    return out


def _figure_size(n_channels: int) -> Tuple[float, float]:
    height = min(max(MIN_HEIGHT, 2.0 + INCH_PER_CHANNEL * n_channels),
                 MAX_HEIGHT)
    return (8.0, height)


def vdedr_plot_task(task: Dict[str, Any]) -> Optional[List[Dict[str, str]]]:
    """Draw the three figures of one selection; returns their paths."""
    try:
        return _plot(task)
    except Exception:
        print(f"[vdedr] plot failed for {task.get('family')} "
              f"{task.get('id_stn')} varno {task.get('varno')}:\n"
              f"{traceback.format_exc()}", flush=True)
        plt.close('all')
        return None


def _plot(task: Dict[str, Any]) -> Optional[List[Dict[str, str]]]:
    stn_sql = task.get('stn_sql', '')
    ctl = _channel_stats(task['files_in'][0], task['region'], task['flag'],
                         stn_sql, task['varno'])
    exp = _channel_stats(task['files_in'][1], task['region'], task['flag'],
                         stn_sql, task['varno'])
    channels = sorted(set(ctl) & set(exp), reverse=True)
    if not channels:
        return None

    name1, name2 = task['names_in']
    period = f"From  {task['datestart']}   to  {task['dateend']}"
    idlev = np.arange(len(channels))
    n1 = np.array([ctl[c]['n'] for c in channels])
    n2 = np.array([exp[c]['n'] for c in channels])

    def col(src, key):
        return np.array([src[c][key] for c in channels], dtype=float)

    with np.errstate(divide='ignore', invalid='ignore'):
        series = {
            'omp_bias': (col(exp, 'avg_omp') - col(ctl, 'avg_omp'),
                         (col(exp, 'avg_omp') - col(exp, 'bcorr'))
                         - (col(ctl, 'avg_omp') - col(ctl, 'bcorr'))),
            'omp_rel': (100.0 * (col(exp, 'std_omp') - col(ctl, 'std_omp'))
                        / col(ctl, 'std_omp'),
                        100.0 * (n2 - n1) / n1),
            'oma_rel': (100.0 * (col(exp, 'std_oma') - col(ctl, 'std_oma'))
                        / col(ctl, 'std_oma'),
                        100.0 * (n2 - n1) / n1),
        }

    conf = {
        'omp_rel': ftest_confidence(_sample_std(col(ctl, 'std_omp'), n1), n1,
                                    _sample_std(col(exp, 'std_omp'), n2), n2),
        'oma_rel': ftest_confidence(_sample_std(col(ctl, 'std_oma'), n1), n1,
                                    _sample_std(col(exp, 'std_oma'), n2), n2),
    }

    out: List[Dict[str, str]] = []
    for metric, (first, second) in series.items():
        typer, legend = METRICS[metric]
        fig, ax = plt.subplots(figsize=_figure_size(len(channels)))
        ax.grid(True, axis='x', linestyle=':', color='#b0b0b0', alpha=0.8)
        # one light band every other channel: with hundreds of rows it is
        # what keeps a value tied to its channel
        for y in idlev[::2]:
            ax.axhspan(y - 0.5, y + 0.5, color=BAND, zorder=0)
        ax.axvline(0.0, color=NOIR, lw=1.0, alpha=0.7, zorder=2)

        signed = metric in ('omp_rel', 'oma_rel')
        ax.plot(first, idlev, linestyle='-', linewidth=1.2, zorder=3,
                color=NEUTRAL_LINE if signed else COULEURS[2],
                label=legend[0])
        if signed:
            # the sign says who wins: negative = experience better
            for values, color, lab in (
                    (np.where(first < 0, first, np.nan), EXP_BETTER,
                     f'{legend[0]} < 0: {name2} better'),
                    (np.where(first > 0, first, np.nan), CTL_BETTER,
                     f'{legend[0]} > 0: {name1} better')):
                ax.plot(values, idlev, linestyle='none', marker='o',
                        markersize=4.5, color=color, zorder=4, label=lab)
        else:
            ax.plot(first, idlev, linestyle='none', marker='o', markersize=4,
                    color=COULEURS[2], zorder=4)
        ax.plot(second, idlev, linestyle='-', marker='p', color=COULEURS[3],
                markersize=4, linewidth=1.2, zorder=3, label=legend[1])

        ax.set_yticks(idlev)
        ax.set_yticklabels([f"{c:g}" for c in channels],
                           fontsize=TICK_FONTSIZE)
        ax.set_ylim(idlev[0] - 0.5, idlev[-1] + 0.5)
        ax.set_ylabel('Channel', color=NOIR, fontsize=16,
                      bbox=dict(facecolor=ROUGE))
        ax.set_xlabel(legend[0] + '   /   ' + legend[1], fontsize=10)

        # F-test confidence of the spread change, as in saska: a square
        # per channel, filled red when the change passes 95 %
        c_vals = conf.get(metric)
        if c_vals is not None:
            show_text = len(channels) <= MAX_CHANNELS_WITH_CONFIDENCE
            for y in idlev:
                if not np.isfinite(c_vals[y]):
                    continue
                face = 'red' if c_vals[y] > MIN_CONFIDENCE else 'white'
                ax.scatter(1.005, y, s=26, marker='s', facecolor=face,
                           edgecolor='black', linewidths=0.4, clip_on=False,
                           zorder=5, transform=ax.get_yaxis_transform())
                if show_text:
                    ax.text(1.03, y, f"{c_vals[y]:.1f}%",
                            fontsize=TICK_FONTSIZE, va='center',
                            fontweight='bold',
                            transform=ax.get_yaxis_transform())
            ax.annotate("F-test", xy=(1.005, 1.0), xycoords='axes fraction',
                        xytext=(0, 5), textcoords='offset points',
                        fontsize=7, fontweight='bold')
            n_sig = int(np.nansum(c_vals > MIN_CONFIDENCE))
            # offsets in points: the figure can be metres tall, so a
            # fraction of the axes would push the header far away
            ax.annotate(f"F-test > {MIN_CONFIDENCE:.0f}%: "
                        f"{n_sig}/{len(channels)}",
                        xy=(1.005, 1.0), xycoords='axes fraction',
                        xytext=(0, 18), textcoords='offset points',
                        fontsize=8, fontweight='bold', color='red')

        if len(channels) <= MAX_CHANNELS_WITH_COUNTS:
            # to the right of the frame, after the confidence column
            shift = 1.16 if c_vals is not None else 1.01
            ax.annotate(f"N {name1}", xy=(shift, 1.0),
                        xycoords='axes fraction', xytext=(0, 5),
                        textcoords='offset points', fontsize=7,
                        fontweight='bold', color=COULEURS[0])
            ax.annotate(f"N {name2}", xy=(shift + 0.06, 1.0),
                        xycoords='axes fraction', xytext=(0, 5),
                        textcoords='offset points', fontsize=7,
                        fontweight='bold', color=COULEURS[1])
            for y in idlev:
                ax.text(shift, y, int(n1[y]), fontsize=TICK_FONTSIZE,
                        color=COULEURS[0], va='center',
                        transform=ax.get_yaxis_transform())
                ax.text(shift + 0.06, y, int(n2[y]), fontsize=TICK_FONTSIZE,
                        color=COULEURS[1], va='center',
                        transform=ax.get_yaxis_transform())

        ax.legend(columnspacing=1.2, handletextpad=0.4, fancybox=False,
                  ncol=2, shadow=False, loc='lower center',
                  bbox_to_anchor=(0.5, 1.0, 0.0, 0.0),
                  bbox_transform=ax.transAxes, prop={'size': 8},
                  frameon=False, borderaxespad=0.6)
        bbox = dict(facecolor=BLEU, boxstyle='round')

        def header(x, y, text, dy):
            ax.annotate(text, xy=(x, y), xycoords='axes fraction',
                        xytext=(0, dy), textcoords='offset points',
                        fontsize=10, bbox=bbox)

        header(-.03, 0.0,
               f"{task['family']} {task['id_stn']}  "
               f"{varno_name(task['varno'])} {typer} {task['flag']}", -52)
        header(.00, 1.0, task['region'], 44)
        header(.25, 1.0, period, 44)
        header(.70, 1.0, f"{name2} - {name1}", 44)

        rel = os.path.join(
            task['family'],
            f"{metric}_{_safe(task['family'])}_"
            f"{_safe(task.get('stn_tag', task['id_stn']))}_"
            f"{_safe(name1)}-{_safe(name2)}_{_safe(task['region'])}_"
            f"{_safe(task['flag'])}_varno{int(task['varno'])}.png")
        path = os.path.join(task['pathwork'], rel)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        fig.savefig(path, format='png', dpi=DPI, bbox_inches='tight')
        plt.close(fig)
        out.append({'metric': metric, 'filename': rel})
    return out


def vdedr_viewer_items(tasks: Sequence[Dict[str, Any]],
                       results: Sequence[Any]) -> List[Dict[str, str]]:
    """One viewer entry per figure actually written."""
    items: List[Dict[str, str]] = []
    for task, res in zip(tasks, results):
        if not res:
            continue
        for fig in res:
            items.append({
                'comparison': task['comparison'],
                'metric': fig['metric'],
                'family': str(task['family']).strip(),
                'region': str(task['region']).strip(),
                'flag': str(task['flag']).strip(),
                'id_stn': str(task['id_stn']).strip(),
                'varno': str(task['varno']).strip(),
                'filename': fig['filename'],
            })
    return items
