from typing import List, Optional, Tuple, Dict
import pandas as pd
import numpy as np
import os
from collections import OrderedDict
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from scipy.stats import mode
from numba import jit

def compute_theta_phi_for_biomarker(biomarker_df, max_attempt = 100, seed = None):
    """get theta and phi parameters for this biomarker using hard k-means
    input: 
        - biomarker_df: a pd.dataframe of a specific biomarker
    output: 
        - a tuple: theta_mean, theta_std, phi_mean, phi_std
    """
    if seed is not None:
        # Set the seed for numpy's random number generator
        rng = np.random.default_rng(seed)
    else:
        rng = np.random

    n_clusters = 2
    measurements = np.array(biomarker_df['measurement']).reshape(-1, 1)
    healthy_df = biomarker_df[biomarker_df['diseased'] == False]

    curr_attempt = 0
    n_init_value = 50
    clustering_setup = KMeans(n_clusters=n_clusters, n_init=n_init_value)
    
    while curr_attempt < max_attempt:
        clustering_result = clustering_setup.fit(measurements)
        predictions = clustering_result.labels_
        cluster_counts = np.bincount(predictions) # array([3, 2])
        
        if len(cluster_counts) == n_clusters and all(c > 1 for c in cluster_counts):
            break 
        curr_attempt += 1
    else:
        print(f"KMeans failed. Try randomizing the predictions")
        predictions = rng.choice([0, 1], size=len(measurements))
        cluster_counts = np.bincount(predictions)
        if len(cluster_counts) != n_clusters or not all(c > 1 for c in cluster_counts):
            raise ValueError(f"KMeans clustering failed to find valid clusters within max_attempt.")
    
    healthy_predictions = predictions[healthy_df.index]
    mode_result = mode(healthy_predictions, keepdims=False).mode
    phi_cluster_idx = mode_result[0] if isinstance(mode_result, np.ndarray) else mode_result
    theta_cluster_idx = 1 - phi_cluster_idx

    # two empty clusters to strore measurements
    clustered_measurements = [[] for _ in range(2)]
    # Store measurements into their cluster
    for i, prediction in enumerate(predictions):
        clustered_measurements[prediction].append(measurements[i][0])
    
    # Calculate means and standard deviations
    theta_mean, theta_std = np.mean(
        clustered_measurements[theta_cluster_idx]), np.std(
            clustered_measurements[theta_cluster_idx])
    phi_mean, phi_std = np.mean(
        clustered_measurements[phi_cluster_idx]), np.std(
            clustered_measurements[phi_cluster_idx])
    
    # Check for invalid values
    if any(np.isnan(v) or v == 0 for v in [theta_std, phi_std, theta_mean, phi_mean]):
        raise ValueError("One of the calculated values is invalid (0 or NaN).")

    return theta_mean, theta_std, phi_mean, phi_std

def get_theta_phi_estimates(
    data: pd.DataFrame,
) -> Dict[str, Dict[str, float]]:
    """
    Obtain theta and phi estimates (mean and standard deviation) for each biomarker.

    Args:
    data (pd.DataFrame): DataFrame containing participant data with columns 'participant', 
        'biomarker', 'measurement', and 'diseased'.
    # biomarkers (List[str]): A list of biomarker names.

    Returns:
    Dict[str, Dict[str, float]]: A dictionary where each key is a biomarker name,
        and each value is another dictionary containing the means and standard deviations 
        for theta and phi of that biomarker, with keys 'theta_mean', 'theta_std', 'phi_mean', 
        and 'phi_std'.
    """
    # empty hashmap of dictionaries to store the estimates
    estimates = {}
    biomarkers = data.biomarker.unique()
    for biomarker in biomarkers:
        # Filter data for the current biomarker
        # reset_index is necessary here because we will use healthy_df.index later
        biomarker_df = data[data['biomarker']
                            == biomarker].reset_index(drop=True)
        theta_mean, theta_std, phi_mean, phi_std = compute_theta_phi_for_biomarker(
            biomarker_df)
        estimates[biomarker] = {
            'theta_mean': theta_mean,
            'theta_std': theta_std,
            'phi_mean': phi_mean,
            'phi_std': phi_std
        }
    return estimates

@jit(nopython=True)
def _compute_likelihood_core(measurements, affected, mus, stds):
    """Core computation function optimized with Numba"""
    likelihood = 1.0 
    for i in range(len(measurements)):
        var = stds[i] * stds[i]
        if var <= 0 or np.isnan(measurements[i]) or np.isnan(mus[i]):
            continue
        likelihood *= np.exp(-(measurements[i] - mus[i])**2 / (2 * var)) / np.sqrt(2 * np.pi * var)
    return likelihood 

# def compute_likelihood(pdata, k_j, theta_phi):
#     """Wrapper function that prepares data for Numba optimization"""
#     measurements = pdata['measurement'].values.astype(np.float64)
#     affected = k_j >= pdata['S_n'].values 

#     biomarkers = pdata['biomarker'].values 
#     mus, stds = np.zeros(len(measurements)), np.zeros(len(measurements))
#     for i, (biomarker, is_affected) in enumerate(zip(biomarkers, affected)):
#         params = theta_phi[biomarker]
#         if is_affected:
#             mus[i] = params['theta_mean']
#             stds[i] = params['theta_std']
#         else:
#             mus[i] = params['phi_mean']
#             stds[i] = params['phi_std']
    
#     return _compute_likelihood_core(measurements, affected, mus, stds)

def compute_likelihood(
    measurements: np.ndarray,
    S_n: np.ndarray,
    biomarkers: np.ndarray,
    k_j: int,
    theta_phi: Dict[str, Dict[str, float]]
) -> float:
    """
    Compute the likelihood for given participant data.

    Args:
        measurements (np.ndarray): Array of measurement values.
        S_n (np.ndarray): Array of stage values (mapped from biomarkers).
        biomarkers (np.ndarray): Array of biomarker names.
        k_j (int): Current stage.
        theta_phi (Dict[str, Dict[str, float]]): Biomarker parameter dictionary.

    Returns:
        float: Likelihood value.
    """
    affected = k_j >= S_n
    mus = np.zeros(len(measurements))
    stds = np.zeros(len(measurements))

    for i, (biomarker, is_affected) in enumerate(zip(biomarkers, affected)):
        params = theta_phi[biomarker]
        if is_affected:
            mus[i] = params['theta_mean']
            stds[i] = params['theta_std']
        else:
            mus[i] = params['phi_mean']
            stds[i] = params['phi_std']

    return utils._compute_likelihood_core(measurements, affected, mus, stds)

def shuffle_order(arr: np.ndarray, n_shuffle: int) -> None:

    """
    Randomly shuffle a specified number of elements in an array.

    Args:
    arr (np.ndarray): The array to shuffle elements in.
    n_shuffle (int): The number of elements to shuffle within the array.
    """
    if n_shuffle > len(arr):
        raise ValueError(
            "n_shuffle cannot be greater than the length of the array")

    # Randomly choose indices to shuffle
    indices = np.random.choice(len(arr), size=n_shuffle, replace=False)

    # Obtain and shuffle the elements at these indices
    selected_elements = arr[indices]
    np.random.shuffle(selected_elements)

    # Place the shuffled elements back into the array
    arr[indices] = selected_elements

def obtain_most_likely_order_dic(all_current_accepted_order_dicts, burn_in, thining):
    """Obtain the most likely order based on all the accepted orders 
    Inputs:
        - all_current_accepted_order_dicts 
        - burn_in
        - thining
    Outputs:
        - a dictionary where key is biomarker and value is the most likely order for that biomarker
        Note that in this dic, the order follows the same order as in 
        data_we_have.biomarker.unique()
    """
    biomarker_stage_probability_df = get_biomarker_stage_probability(
        all_current_accepted_order_dicts, burn_in, thining)
    od = OrderedDict()
    assigned_stages = set()

    for i, biomarker in enumerate(biomarker_stage_probability_df.index):
        # probability array for that biomarker
        prob_arr = np.array(biomarker_stage_probability_df.iloc[i, :])

        # Sort indices of probabilities in descending order
        sorted_indices = np.argsort(prob_arr)[::-1] + 1

        for stage in sorted_indices:
            if stage not in assigned_stages:
                od[biomarker] = int(stage)
                assigned_stages.add(stage)
                break
        else:
            raise ValueError(
                f"Could not assign a unique stage for biomarker {biomarker}.")
    return od

def get_biomarker_stage_probability(all_current_accepted_order_dicts, burn_in, thining):
    """filter through all_dicts using burn_in and thining 
    and for each biomarker, get probability of being in each possible stage

    Input:
        - all_current_accepted_order_dicts 
        - burn_in
        - thinning
    Output:
        - dff: a pandas dataframe where index is biomarker name, each col is each stage
        and each cell is the probability of that biomarker indicating that stage

        Note that in dff, its index follows the same order as data_we_have.biomarker.unique()
    """
    df = pd.DataFrame(all_current_accepted_order_dicts)
    df = df[(df.index > burn_in) & (df.index % thining == 0)]
    # Create an empty list to hold dictionaries
    dict_list = []

    # biomarkers are in the same order as data_we_have.biomarker.unique()
    biomarkers = np.array(df.columns)

    # iterate through biomarkers
    for biomarker in biomarkers:
        dic = {"biomarker": biomarker}
        # get the frequency of biomarkers
        # value_counts will generate a Series where index is each cell's value
        # and the value is the frequency of that value
        stage_counts = df[biomarker].value_counts()
        # for each stage
        # note that df.shape[1] should be equal to num_biomarkers
        for i in range(1, df.shape[1] + 1):
            # get stage:prabability
            dic[i] = stage_counts.get(i, 0)/len(df)
        dict_list.append(dic)

    dff = pd.DataFrame(dict_list)
    dff.set_index(dff.columns[0], inplace=True)
    return dff

def save_heatmap(all_dicts, burn_in, thining, folder_name, file_name, title):
    # Check if the directory exists
    if not os.path.exists(folder_name):
        # Create the directory if it does not exist
        os.makedirs(folder_name)
    biomarker_stage_probability_df = get_biomarker_stage_probability(
        all_dicts, burn_in, thining)
    sns.heatmap(biomarker_stage_probability_df,
                annot=True, cmap="Greys", linewidths=.5,
                cbar_kws={'label': 'Probability'},
                fmt=".1f",
                # vmin=0, vmax=1,
                )
    plt.xlabel('Stage')
    plt.ylabel('Biomarker')
    plt.title(title)
    plt.tight_layout()
    plt.savefig(f"{folder_name}/{file_name}.png")
    # plt.savefig(f'{file_name}.pdf')
    plt.close()
