"""Charts for the dataset and for the model results.

Dataset charts
    :func:`plot_sentiment_distribution` - class balance
    :func:`plot_review_length_distribution` - words per review, raw vs cleaned
    :func:`plot_top_tokens` - most frequent tokens per sentiment

Model charts
    :func:`plot_model_comparison` - all four models across all four metrics
    :func:`plot_confusion_matrix` - where the best model confuses classes
    :func:`plot_per_class_metrics` - precision/recall/F1 per sentiment
"""

from collections import Counter

import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, precision_recall_fscore_support

from config.constants import LABEL_COLUMN, TEXT_COLUMN
from utils.chart_style import (
    ACCENT,
    INK,
    INK_MUTED,
    INK_SECONDARY,
    SENTIMENT_COLORS,
    SERIES,
    apply_chart_style,
    finish,
    sequential_cmap,
    strip_spines,
)

#: Plot classes in polarity order, not alphabetically.
LABEL_ORDER = ["Negative", "Neutral", "Positive"]


def _ordered_labels(values):
    """Return LABEL_ORDER filtered to the labels actually present."""
    present = set(values)
    return [label for label in LABEL_ORDER if label in present]


def plot_sentiment_distribution(df, label_column=LABEL_COLUMN,
                                title='Distribution of Sentiments',
                                save_path=None, show=True):
    """Bar chart of the row count per sentiment. Counts start at zero, so bar length is honest."""
    apply_chart_style()

    counts = df[label_column].value_counts()
    labels = _ordered_labels(counts.index)
    values = [counts[label] for label in labels]

    fig, ax = plt.subplots(figsize=(7, 4.5))
    bars = ax.bar(labels, values,
                  color=[SENTIMENT_COLORS[label] for label in labels],
                  width=0.6)

    # only three bars, so a direct label on each is still selective
    for bar, value in zip(bars, values):
        ax.annotate(f'{value:,}',
                    (bar.get_x() + bar.get_width() / 2, bar.get_height()),
                    textcoords='offset points', xytext=(0, 4),
                    ha='center', fontsize=10, color=INK_SECONDARY)

    ax.set_xlabel('')
    ax.set_ylabel('Reviews')
    ax.set_title(title)
    ax.set_ylim(0, max(values) * 1.12)
    ax.grid(axis='x', visible=False)
    strip_spines(ax)

    return finish(fig, save_path, show)


def plot_review_length_distribution(df, text_column=TEXT_COLUMN,
                                    compare=None, max_words=60,
                                    title='Review length',
                                    save_path=None, show=True):
    """Histogram of words per review; optionally overlay a second corpus."""
    apply_chart_style()

    def lengths(frame):
        return frame[text_column].dropna().astype(str).str.split().str.len()

    fig, ax = plt.subplots(figsize=(7.5, 4.5))
    bins = np.arange(0, max_words + 2)

    if compare is not None:
        # outlines, not translucent fills: overlapping alpha would invent a third colour that is in neither series
        ax.hist(lengths(compare).clip(upper=max_words), bins=bins,
                histtype='step', linewidth=2, color=SERIES[1],
                label='Before cleaning')
        ax.hist(lengths(df).clip(upper=max_words), bins=bins,
                histtype='step', linewidth=2, color=SERIES[0],
                label='After cleaning')
        ax.legend(loc='upper right')
    else:
        ax.hist(lengths(df).clip(upper=max_words), bins=bins, color=ACCENT)

    median = float(lengths(df).median())
    ax.axvline(median, color=INK_MUTED, linewidth=1)
    ax.annotate(f'median {median:.0f} words', (median, ax.get_ylim()[1]),
                textcoords='offset points', xytext=(6, -12),
                fontsize=9, color=INK_SECONDARY)

    ax.set_xlabel(f'Words per review (clipped at {max_words})')
    ax.set_ylabel('Reviews')
    ax.set_title(title)
    ax.grid(axis='x', visible=False)
    strip_spines(ax)

    return finish(fig, save_path, show)


def plot_top_tokens(df, text_column=TEXT_COLUMN, label_column=LABEL_COLUMN,
                    top_n=12, title='Most frequent tokens by sentiment',
                    save_path=None, show=True):
    """Small multiples: the ``top_n`` most frequent cleaned tokens per class.
    Run this on the cleaned text - it shows what the model actually sees.
    """
    apply_chart_style()

    labels = _ordered_labels(df[label_column].unique())
    fig, axes = plt.subplots(1, len(labels), figsize=(4.2 * len(labels), 5.2),
                             sharex=False)
    if len(labels) == 1:
        axes = [axes]

    for ax, label in zip(axes, labels):
        texts = df.loc[df[label_column] == label, text_column].dropna().astype(str)
        counter = Counter(token for text in texts for token in text.split())
        common = counter.most_common(top_n)[::-1]
        tokens = [token for token, _ in common]
        counts = [count for _, count in common]

        ax.barh(range(len(tokens)), counts,
                color=SENTIMENT_COLORS[label], height=0.65)
        ax.set_yticks(range(len(tokens)))
        ax.set_yticklabels(tokens, fontsize=10)
        ax.set_title(label, color=SENTIMENT_COLORS[label])
        ax.set_xlabel('Occurrences')
        ax.grid(axis='y', visible=False)
        strip_spines(ax)

    fig.suptitle(title, fontsize=13, color=INK, y=1.01)
    return finish(fig, save_path, show)


def plot_model_comparison(results_df, metrics=('F1 Score', 'Precision',
                                               'Recall', 'Accuracy'),
                          title='Model comparison',
                          save_path=None, show=True):
    """Small multiples of dot plots - one facet per metric, models on the y-axis."""
    apply_chart_style()

    models = list(results_df['Model'])
    y = np.arange(len(models))

    fig, axes = plt.subplots(1, len(metrics), figsize=(3.4 * len(metrics), 3.6),
                             sharey=True)
    if len(metrics) == 1:
        axes = [axes]

    for ax, metric in zip(axes, metrics):
        values = results_df[metric].to_numpy(dtype=float)
        best = int(np.argmax(values))
        # ties are real - accent every model that reaches the maximum
        colors = [ACCENT if value == values.max() else INK_MUTED
                  for value in values]

        ax.hlines(y, values.min(), values, color='#e1e0d9', linewidth=1.2)
        ax.scatter(values, y, s=70, color=colors, zorder=3)

        # label only the winner - a number on every dot would be noise
        ax.annotate(f'{values[best]:.4f}', (values[best], y[best]),
                    textcoords='offset points', xytext=(0, 11),
                    ha='center', fontsize=9, color=INK)

        span = max(values.max() - values.min(), 1e-3)
        ax.set_xlim(values.min() - span * 0.45, values.max() + span * 0.45)
        ax.set_title(metric, fontsize=11)
        ax.set_yticks(y)
        ax.set_yticklabels(models, fontsize=10)
        ax.grid(axis='y', visible=False)
        ax.tick_params(axis='x', labelsize=8)
        ax.tick_params(axis='y', length=0)   # labels only; no floating dashes
        strip_spines(ax, keep=('bottom',))

    fig.suptitle(title, fontsize=13, color=INK, y=1.04)
    return finish(fig, save_path, show)


def plot_confusion_matrix(y_test, y_pred, normalize=True,
                          title='Confusion matrix',
                          save_path=None, show=True):
    """Heatmap of true vs predicted labels for one model."""
    apply_chart_style()

    labels = _ordered_labels(set(y_test) | set(y_pred))
    matrix = confusion_matrix(y_test, y_pred, labels=labels)
    shown = (matrix / matrix.sum(axis=1, keepdims=True)) if normalize else matrix

    fig, ax = plt.subplots(figsize=(5.6, 5))
    image = ax.imshow(shown, cmap=sequential_cmap(), vmin=0,
                      vmax=shown.max())

    threshold = shown.max() * 0.6
    for i in range(len(labels)):
        for j in range(len(labels)):
            cell = f'{shown[i, j]:.1%}' if normalize else f'{matrix[i, j]:,}'
            sub = f'\n{matrix[i, j]:,}' if normalize else ''
            ax.text(j, i, cell + sub, ha='center', va='center', fontsize=10,
                    color='#ffffff' if shown[i, j] > threshold else INK)

    ax.set_xticks(range(len(labels)), labels)
    ax.set_yticks(range(len(labels)), labels)
    ax.set_xlabel('Predicted')
    ax.set_ylabel('Actual')
    ax.set_title(title)
    ax.grid(False)
    strip_spines(ax, keep=())

    bar = fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
    bar.outline.set_visible(False)
    bar.ax.tick_params(labelsize=8, color=INK_MUTED)

    return finish(fig, save_path, show)


def plot_per_class_metrics(y_test, y_pred, title='Per-class performance',
                           save_path=None, show=True):
    """Precision / recall / F1 for each sentiment class, as a grouped dot plot."""
    apply_chart_style()

    labels = _ordered_labels(set(y_test) | set(y_pred))
    precision, recall, f1, support = precision_recall_fscore_support(
        y_test, y_pred, labels=labels, zero_division=0)

    series = [('Precision', precision, SERIES[0]),
              ('Recall', recall, SERIES[1]),
              ('F1', f1, SERIES[2])]

    fig, ax = plt.subplots(figsize=(7.5, 4.4))
    y = np.arange(len(labels))
    offsets = (0.22, 0.0, -0.22)   # top-to-bottom matches the legend order

    for (name, values, color), offset in zip(series, offsets):
        ax.scatter(values, y + offset, s=70, color=color, label=name, zorder=3)
        for value, position in zip(values, y + offset):
            ax.annotate(f'{value:.3f}', (value, position),
                        textcoords='offset points', xytext=(9, -3),
                        fontsize=8.5, color=INK_SECONDARY)

    ax.set_yticks(y)
    ax.set_yticklabels([f'{label}\n(n={count:,})'
                        for label, count in zip(labels, support)], fontsize=10)
    ax.set_xlabel('Score')
    ax.set_title(title)

    lowest = float(min(precision.min(), recall.min(), f1.min()))
    ax.set_xlim(max(0.0, lowest - 0.06), 1.02)
    ax.grid(axis='y', visible=False)

    # legend below the plot - inside the axes it collides with the bottom row
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.16), ncol=3)
    strip_spines(ax, keep=('bottom',))

    return finish(fig, save_path, show)