"""
plot_segment_test.py — Funzioni di visualizzazione per i Stress Test di mercato

Fornisce la funzione plot_choice_probabilities() usata da tutti i test_DROP_*.py
e test_filter_*.py per visualizzare come cambiano le probabilità di scelta dei
clienti quando una parte del catalogo viene rimossa.

Il grafico mostra:
  - Una linea colorata per ogni segmento (loyal_A/B/C/D, price-sensitive, quality-sensitive)
  - Una linea nera spessa per la probabilità aggregata del Latent Class Model
  - Sull'asse X: tutti i prodotti rimasti dopo il filtro/drop
  - Sull'asse Y: probabilità di scelta (in percentuale)

Come usare:
  from plot_segment_test import plot_choice_probabilities

  # res è il risultato di sim_*StressTest(), che ritorna [df_info, P_avg, P_seg]
  res = test.sim_MarcaDropStressTest(MNLChoiceModel, LatentChoiceModel)
  plot_choice_probabilities(
      res,
      test_label="Drop Marca A",
      segment_names=["Loyal A", "Loyal B", "Loyal C", "Loyal D", "Price", "Quality"]
  )

"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from matplotlib.lines import Line2D

seg_color=["#3B82F6","#16A34A","#EF4444","#F59E0B","#8B5CF6","#06B6D4","#EC4899","#84CC16","#F97316","#14B8A6","#6366F1","#D97706",]
latent_color="#111827"

def _to_array(x):
    """funzione richiamata da plot_choice_probabilities per convertire il dataset in un array di numpy"""
    if isinstance(x, pd.DataFrame): return x.values
    return np.asarray(x)

def _build_product_labels(df):
    """funzione richiamata da plot_choice_probabilities per inserire un'etichetta ad ogni prodotto e alla no_buy option"""
    labels = []
    i=0
    for _, _ in df.iterrows():
        i+=1
        labels.append(f"P_{i}")
    labels.append(f"No-Buy")
    return labels

def plot_choice_probabilities(res,test_label="Stress Test",segment_names=None,figsize=None,):
    products_df= res[0].reset_index(drop=True) #serve per prendere le etichette dei prodotti
    probs_avg= _to_array(res[1]).flatten() #serve per le probabilità associate ai prodotti dal mercato 
    probs_seg= _to_array(res[2]) #serve per le probabilità associate ai prodotti da ogni segmento
    if probs_seg.ndim==1:
        probs_seg=probs_seg.reshape(-1, len(probs_avg))

    n_seg=probs_seg.shape[0]
    n_prod=probs_seg.shape[1]

    if segment_names is None:
        segment_names=[f"Seg {i+1}" for i in range(n_seg)]


    labels=_build_product_labels(products_df)
    x=np.arange(n_prod)

    #creiamo il plot con x: prodotti, y:valore di probabilità
    if figsize is None:
        figsize=(max(28, n_prod * 0.30), 8)

    fig,ax= plt.subplots(figsize=figsize, facecolor="white")
    ax.set_facecolor("#F8FAFC")

    for i in range(0, n_prod, 2):
        ax.axvspan(i - 0.5, i + 0.5, color="#E2E8F0", alpha=0.45, zorder=0)

    for i in range(n_seg):
        color = seg_color[i % len(seg_color)]
        ax.plot(x, probs_seg[i],color=color,linewidth=1.1,alpha=0.65,marker="o",markersize=2.5,markeredgewidth=0,zorder=2, label=segment_names[i])

    ax.plot(x, probs_avg,color=latent_color,linewidth=2.6,alpha=0.95,marker="D",markersize=4,markeredgewidth=0,zorder=5, label="Modello Latent Class")

    ax.set_xticks(x)
    ax.set_xticklabels(labels,fontsize=5.2,rotation=90,linespacing=1.25, color="#374151")
    ax.xaxis.set_tick_params(length=3, width=0.5)

    ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v:.2%}" if probs_avg.max()>0.005 else f"{v:.4%}"))
    ax.tick_params(axis="y",labelsize=8,colors="#6B7280")

    ax.spines[["top","right"]].set_visible(False)
    ax.spines[["left","bottom"]].set_color("#CBD5E1")
    ax.grid(axis="y",color="#E2E8F0",linewidth=0.7,zorder=1)
    ax.set_axisbelow(True)
    ax.set_xlim(-0.8,n_prod-0.2)
    ax.set_ylabel("Probabilità di scelta",fontsize=9,color="#374151",labelpad=8)

    fig.suptitle(f"Distribuzione probabilità di scelta {test_label}",fontsize=13,fontweight="bold",color="#111827",y=1.01)

    handles = [Line2D([0],[0], color=seg_color[i % len(seg_color)],linewidth=1.5,marker="o",markersize=4, label=nm) for i, nm in enumerate(segment_names)] + [Line2D([0],[0], color=latent_color, linewidth=2.5, marker="D", markersize=5,label="Modello Latent Class")]
    ax.legend(handles=handles,fontsize=7.5,title="Segmenti",title_fontsize=8,loc="upper right", frameon=True, framealpha=0.92,edgecolor="#CBD5E1",ncol=1,bbox_to_anchor=(1.0, 1.0),)

    plt.tight_layout()
    plt.show()
    return fig

