import pandas as pd
import matplotlib.pyplot as plt
from riskreg.phi import phi
from sklearn.mixture import GaussianMixture
from scipy.stats import gaussian_kde
import numpy as np
from riskreg.phi_ctrl_pts import phi_ctrl_pts
import pickle
from datetime import datetime
import os
import matplotlib.cm as cm
import matplotlib.colors as colors
from matplotlib.collections import LineCollection
import numpy as np

class PhiRelevance:
    def __init__(self, ctrl_pts, method):
        self.ctrl_pts = ctrl_pts
        self.method = method

    def __call__(self, v):
        # wrap the single value in a Series so phi() can reset its index
        series = pd.Series([v])
        return phi(series, ctrl_pts=self.ctrl_pts, method=self.method)[0]

class GMMRelevance:
    def __init__(self, gmm):
        self.gmm = gmm
    def __call__(self, v):
        # return density
        return float(np.exp(self.gmm.score_samples([[v]])))

class KDERelevance:
    def __init__(self, data):
        self.kde = gaussian_kde(np.array(data))
    def __call__(self, v):
        return float(self.kde(v))

class HistRelevance:
    def __init__(self, centers, phis):
        self.centers = np.array(centers)
        self.phis = np.array(phis)
    def __call__(self, v):
        idx = np.argmin(np.abs(self.centers - v))
        return float(self.phis[idx])

class SpectralRelevance:
    def __init__(self, y, phi_scores):
        self.y = np.asarray(y)
        self.phi_scores = np.asarray(phi_scores)

    def __call__(self, v):
        idx = np.argmin(np.abs(self.y - v))
        return float(self.phi_scores[idx])

# === Load dataset === #
def returnPhiRelevance(path):
    datasets = os.listdir(path)
    print("\nAvailable datasets found:")
    for i, dataset in enumerate(datasets, 1):
        print(f"{i}. {dataset}")
    choice = int(input("\nNumber the dataset: "))
    # specify path to dataset or dataset directory
    path = f"{path}/{datasets[choice - 1]}"
    name = path.split("/")[-1].split(".")[0]
    df = pd.read_csv(path)
    target = df.columns[-1]
    y = df.iloc[:, -1]
    if(target == "TotalPayBenefits"):
        y = np.log1p(y)
        df["TotalPayBenefits"] = y
        df = df[df["TotalPayBenefits"] > 0]
        y = df.iloc[:, -1]

    # === User choice === #
    print("\nChoose phi relevance method:")
    print("1. default (auto boxplot-based)")
    print("2. manual (user-specified control points)")
    print("3. gmm (Gaussian Mixture Model)")
    print("4. kde (Kernel Density Estimation)")
    print("5. hist (Histogram-based)")
    print("6. spectral (Spectral Density Estimation)")

    choice = input("\nEnter method name or number (1-6): ").strip().lower()

    method_map = {
        "1": "default",
        "2": "manual",
        "3": "gmm",
        "4": "kde",
        "5": "hist",
        "6": "spectral"
    }

    method = method_map.get(choice, choice)

    # === Compute phi values === #
    if method == "default":
        ctrl_pts = phi_ctrl_pts(y, coef=1 ,method="auto")
        y_phi = phi(y, ctrl_pts=ctrl_pts, method="default")

    elif method == "manual":
        print("\nEnter manual control points as list of [x, y, m] triples.")
        print("Example: [[100000, 1, 0], [200000, 0, 0], [400000, 1, 0]]")
        user_input = input("Paste your list here: ").strip()
        try:
            import ast
            custom_pts = ast.literal_eval(user_input)
            ctrl_pts = phi_ctrl_pts(y, method="manual", ctrl_pts=custom_pts)
            y_phi = phi(y, ctrl_pts=ctrl_pts, method="default")
        except Exception as e:
            print(f"Error parsing manual control points: {e}")
            exit(1)
    elif method == "hist":
        bin_centers_phi, phi_at_mid, y_phi = phi(y, method=method)
    else:
        y_phi = phi(y, method=method)

    # === Add to DataFrame === #
    df["phi"] = y_phi

    # === Save to CSV === #
    out_file = f"{name}_phi_{method}.csv"
    df.to_csv(f"relevance_csv/{out_file}", index=False)
    print(f"\nPhi values computed using '{method}' method and saved to '{out_file}'.\n")

    # === Interpolation Prep === #
    y_np = np.array(y)
    phi_np = np.array(y_phi)
    sorted_idx = np.argsort(y_np)
    x_sorted = y_np[sorted_idx]
    phi_sorted = phi_np[sorted_idx]

    x_unique, idx_unique = np.unique(x_sorted, return_index=True)
    phi_unique = phi_sorted[idx_unique]
    slopes = np.gradient(phi_unique, x_unique)
    triplets = list(zip(x_unique, phi_unique, slopes))
    flattened_ctrl_pts = [val for triplet in triplets for val in triplet]
    ctrl_pts_from_method = {
        "ctrl_pts": flattened_ctrl_pts,
        "num_pts": len(triplets)
    }

    # Interpolated phi using Hermite spline (same logic)
    y_series = pd.Series(x_unique)
    phi_interp = phi(y_series, ctrl_pts=ctrl_pts_from_method, method="default")

    # === Combined Plot === #
    fig, axs = plt.subplots(1, 3, figsize=(18, 5), constrained_layout=True)
    xs = np.linspace(min(y), max(y), 500)

    # Plot 1: Density
    if method == "gmm":
        y_reshaped = np.array(y).reshape(-1, 1)
        gmm = GaussianMixture(n_components=3, random_state=0).fit(y_reshaped)
        log_probs = gmm.score_samples(xs.reshape(-1, 1))
        axs[0].plot(xs, np.exp(log_probs), color='purple')
        axs[0].set_title("GMM Estimated Density")

    elif method == "kde":
        kde = gaussian_kde(y)
        axs[0].plot(xs, kde(xs), color='green')
        axs[0].set_title("KDE Estimated Density")

    elif method == "hist":
        counts, bin_edges = np.histogram(y, bins=20, density=True)
        bin_centers = 0.5 * (bin_edges[1:] + bin_edges[:-1])
        axs[0].plot(bin_centers, counts, marker='o', color='dodgerblue')
        axs[0].set_title("Histogram Density")

    else:
        kde = gaussian_kde(y)
        axs[0].plot(xs, kde(xs), color='gray')
        axs[0].set_title("KDE (Generic) Density")

    axs[0].set_xlabel(f"{target}")
    axs[0].set_ylabel("Density")
    axs[0].grid(True)

    # Plot 2: Raw phi
    if method == "hist":
        cmap = cm.viridis
        norm = colors.Normalize(vmin=min(phi_at_mid), vmax=max(phi_at_mid))
        # Create colored line segments
        points = np.array([bin_centers_phi, phi_at_mid]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)

        lc = LineCollection(segments, cmap=cmap, norm=norm)
        lc.set_array(phi_at_mid)       # color by phi values
        lc.set_linewidth(2)
        line = axs[1].add_collection(lc)

        # Overlay midpoints as visible markers
        axs[1].scatter(bin_centers_phi, phi_at_mid,
                    c=phi_at_mid, cmap=cmap, norm=norm,
                    edgecolor="k", s=50, zorder=3)

        axs[1].set_title(f"Raw ϕ at Midpoints - {method}")
        axs[1].set_xlabel(f"{target}")
        axs[1].set_ylabel("ϕ Relevance")
        axs[1].grid(True)
        # cmap = cm.viridis   # <-- change to the one you want
        # norm = colors.Normalize(vmin=min(phi_at_mid), vmax=max(phi_at_mid))
        # axs[1].plot(bin_centers_phi, phi_at_mid, color="black", alpha=0.5, linewidth=1)

        # # Overlay scatter points with colormap
        # sc = axs[1].scatter(bin_centers_phi, phi_at_mid,
        #                     c=phi_at_mid, cmap=cmap, norm=norm,
        #                     edgecolor="k", s=40, zorder=3)

        # axs[1].set_title(f"Raw ϕ at Midpoints - {method}")
        # axs[1].set_xlabel(f"{target}")
        # axs[1].set_ylabel("ϕ Relevance")
        # axs[1].grid(True)

        # axs[1].plot(bin_centers_phi, phi_at_mid, marker='o')
        # axs[1].set_title(f"Raw ϕ at Midpoints - {method}")
        # axs[1].grid(True)
    else:
        axs[1].scatter(df[f"{target}"], df["phi"], c=df["phi"], alpha=0.7, s=10)
        axs[1].set_title(f"Raw ϕ Scores - Method: {method}")
        axs[1].set_xlabel(f"{target}")
        axs[1].set_ylabel("ϕ Relevance")
        axs[1].grid(True)

    # Plot 3: Interpolated phi
    if method == "hist":
        phi_interp_vals = phi(pd.Series(bin_centers_phi), ctrl_pts=ctrl_pts_from_method, method="default")
        axs[2].plot(bin_centers_phi, phi_interp_vals, color='darkgreen', marker='o')
        axs[2].set_title(f"Interpolated ϕ Curve - {method}")
        axs[2].grid(True)
    else:
        axs[2].plot(x_unique, phi_interp, color='darkgreen', linewidth=2)
        axs[2].set_title(f"Interpolated ϕ Curve - {method}")
        axs[2].set_xlabel(f"{target}")
        axs[2].set_ylabel("ϕ Relevance")
        axs[2].grid(True)

    # Match x-axis range
    x_min, x_max = min(y), max(y)
    for ax in axs:
        ax.set_xlim(x_min, x_max)

    # Save combined plot
    plt.savefig(f"relevance_plots/{name}_phi_combined_{method}.png")
    plt.close()
    print(f"Combined plot saved to relevance_plots/{name}_phi_combined_{method}.png\n")

    # pickle logic ===
    # Precompute histogram centers & φ-midpoints if needed
    if method == "hist":
        bin_indices = np.digitize(y, bin_edges, right=False)
        bin_indices = np.asarray(bin_indices).flatten().astype(int)
        y_phi = np.asarray(y_phi).flatten() 
        phi_mid = [
            float(y_phi[bin_indices == i].mean()) if np.any(bin_indices == i) else 0.0
            for i in range(1, len(bin_edges))
        ]

    # Instantiate module-level relevance callable
    if method in ("default", "manual"):
        relevance_func = PhiRelevance(ctrl_pts_from_method, method)
    elif method == "gmm":
        relevance_func = GMMRelevance(gmm)
    elif method == "kde":
        relevance_func = KDERelevance(y)
    elif method == "hist":
        relevance_func = HistRelevance(bin_centers, phi_mid)
    elif method == "spectral":
        relevance_func = SpectralRelevance(y, y_phi)  
    else:
        raise ValueError(f"Unsupported φ-method: {method}")


    # Prepare pickle payload
    pickle_data = {
        "Data":        name + ".csv",
        "method":      method,
        "target_name": df.columns[-1],
        "relevance":   relevance_func,
        "timestamp":   datetime.now().isoformat()
    }

    out_path = f"phi_pickles/{name}_phi_model_{method}.pkl"
    with open(out_path, "wb") as f:
        pickle.dump(pickle_data, f, protocol=pickle.HIGHEST_PROTOCOL)
    print(f"Pickle saved to {out_path}")

    # === Print Top 10 Rare Points === #
    top_phi = df.sort_values("phi", ascending=False).head(10)
    print("\nTop 10 most relevant points (phi):\n")
    for _, row in top_phi.iterrows():
        print(f"{target}: {int(row[f'{target}'])}, Phi: {row['phi']:.3f}")


# main function
if __name__=="__main__":
    # specify path to dataset or dataset directory
    path = f"/Users/bhavneetsingh/Desktop/RiskReg-main/data"
    returnPhiRelevance(path)