import pandas as pd
import numpy as np
import re
from sklearn.preprocessing import PowerTransformer

# 1. Load Data
antigen1 = pd.read_excel('../xmap_biomarkers/data/02a.AP0211_GLA02_SBA01_Antigen_list.xlsx')
antigen2 = pd.read_excel('../xmap_biomarkers/data/02b. AP0211 GLA02_SBA02_Antigen_list.xlsx')
data1 = pd.read_excel('../xmap_biomarkers/data/12a. AP0211 GLA02 SBA01_Data Intensity.xlsx')
data2 = pd.read_excel('../xmap_biomarkers/data/12b. AP0211 GLA02 SBA02_Data_Intensity.xlsx')
layout = pd.read_excel('../xmap_biomarkers/data/layout.xlsx')

datasets = [data1, data2]
antigens = [antigen1, antigen2]
controls = ["Anti-human IgG", "EBNA1", "Bare-bead", "His6ABP"]

# Helper functions
def handle_background(df):
    # Empty sample is "EMPTY-0001"
    empty_df = df[df['Internal LIMS ID'] == 'EMPTY-0001']
    if empty_df.empty:
        return df
    
    # Calculate background noise
    # We only care about Analyte columns
    analyte_cols = [c for c in df.columns if c.startswith('Analyte')]
    empty_means = empty_df[analyte_cols].mean()
    median_mean = empty_means.median()
    sd_mean = empty_means.std()
    
    # subset of empty_means < median + sd
    inset = empty_means[empty_means < (median_mean + sd_mean)]
    if len(inset) < 0.95 * len(empty_means):
        calset = empty_means
    else:
        calset = inset
        
    cutoff = calset.max() + calset.std()
    
    # Keep columns where at least one sample is > cutoff
    bgmap = df[analyte_cols] > cutoff
    keep_cols = bgmap.sum() > 0
    keep_analyte_cols = [c for c in analyte_cols if keep_cols[c]]
    
    return df[['Internal LIMS ID', 'Original ID'] + keep_analyte_cols]

def emptyadjust(df):
    empty_df = df[df['Internal LIMS ID'] == 'EMPTY-0001']
    if empty_df.empty:
        return df
    
    analyte_cols = [c for c in df.columns if c.startswith('Analyte')]
    emptyvector = empty_df[analyte_cols].mean()
    
    df = df[(df['Internal LIMS ID'] != 'EMPTY-0001') & (df['Internal LIMS ID'] != 'MIX_2-0029')].copy()
    
    # Subtract empty vector
    df[analyte_cols] = df[analyte_cols].sub(emptyvector, axis=1)
    
    # set <0 to 0, then add 1
    df[analyte_cols] = df[analyte_cols].clip(lower=0) + 1
    return df

def compress_duplicates(df, layout_df):
    # Layout merges based on Tube label having a hyphen
    # Find base names before hyphen
    tube_labels = layout_df['Tube label'].dropna()
    hyphen_labels = tube_labels[tube_labels.str.contains('-')]
    
    base_names = []
    for lbl in hyphen_labels:
        base = lbl.split('-')[0]
        if base not in base_names:
            base_names.append(base)
            
    analyte_cols = [c for c in df.columns if c.startswith('Analyte')]
            
    for base in base_names:
        # Find sample ids in layout matching base$ or base-
        pattern = f"^{base}$|^{base}-"
        matching_layout = layout_df[layout_df['Tube label'].str.contains(pattern, regex=True, na=False)]
        mergerows = matching_layout['Sample id_LIMS'].dropna().tolist()
        
        if not mergerows:
            continue
            
        # Find these sample ids in df
        regex_pattern = '|'.join(mergerows)
        matching_idx = df['Internal LIMS ID'].str.contains(regex_pattern, regex=True, na=False)
        
        if matching_idx.sum() > 0:
            # calculate mean
            mean_vals = df.loc[matching_idx, analyte_cols].mean()
            # replace first occurrence
            first_idx = df[matching_idx].index[0]
            df.loc[first_idx, analyte_cols] = mean_vals
            # drop others
            drop_idx = df[matching_idx].index[1:]
            df = df.drop(drop_idx)
            
    return df

def set_colname_adapter(df, antigen_df):
    mapping = {}
    for col in df.columns:
        if col.startswith('Analyte'):
            num = int(col.split(' ')[1])
            match = antigen_df[antigen_df['BeadID (Analyte)'] == num]
            if not match.empty:
                antigen_name = match.iloc[0]['Antigen name']
                mapping[col] = antigen_name
    df = df.rename(columns=mapping)
    # Remove controls
    drop_cols = [c for c in df.columns if c in controls]
    df = df.drop(columns=drop_cols)
    return df

# Apply stage 1
processed_datasets = []
for i in range(2):
    df = datasets[i]
    df = handle_background(df)
    df = emptyadjust(df)
    df = compress_duplicates(df, layout)
    df = set_colname_adapter(df, antigens[i])
    # Extract Group based on Original ID (GC -> 1, HD -> 0, else NaN)
    df['group'] = df['Original ID'].apply(lambda x: 1 if 'GC' in str(x) else (0 if 'HD' in str(x) else np.nan))
    processed_datasets.append(df)

# Stage 2: Merge down
df1, df2 = processed_datasets
# Align columns: Internal LIMS ID, Original ID, group
common_keys = ['Internal LIMS ID', 'Original ID', 'group']
all_cols = set(df1.columns).union(set(df2.columns))
analyte_cols_all = list(all_cols - set(common_keys))

# Since df1 and df2 have same rows, we can merge on Internal LIMS ID
merged = pd.merge(df1, df2, on=['Internal LIMS ID', 'Original ID', 'group'], how='outer', suffixes=('_1', '_2'))

# Average common columns
final_cols = {}
for col in analyte_cols_all:
    if col + '_1' in merged.columns and col + '_2' in merged.columns:
        merged[col] = merged[[col + '_1', col + '_2']].mean(axis=1)
        merged = merged.drop(columns=[col + '_1', col + '_2'])
    elif col + '_1' in merged.columns:
        merged = merged.rename(columns={col + '_1': col})
    elif col + '_2' in merged.columns:
        merged = merged.rename(columns={col + '_2': col})

# Drop rows with NaN group
merged = merged.dropna(subset=['group'])
merged = merged.reset_index(drop=True)

# Separate features and target
X = merged.drop(columns=common_keys)
y = merged['group']
meta = merged[['Internal LIMS ID', 'Original ID']]

# Box-Cox transformation + Standardization via PowerTransformer
pt = PowerTransformer(method='box-cox', standardize=True)
# Ensure strictly positive values for Box-Cox
min_val = X.min().min()
if min_val <= 0:
    X = X - min_val + 1e-5

X_transformed = pt.fit_transform(X)
X_transformed_df = pd.DataFrame(X_transformed, columns=X.columns)

# Final dataset
final_df = pd.concat([meta, y.astype(int), X_transformed_df], axis=1)

# Now, create the slices!
# 1) Significant 7
sig_7 = ['HPRA000767', 'HPRA034083', 'HPRA006876', 'HPRA019035', 'HPRA003490', 'HPRA022019', 'HPRA017192']
# Filter out any that might have been dropped during background removal
sig_7_present = [c for c in sig_7 if c in final_df.columns]
slice_1 = final_df[common_keys + sig_7_present]
slice_1.to_csv('slice_1_significant_7.csv', index=False)

# 2) Significant 7 + 33 random proteins = 40 total
np.random.seed(42)
other_proteins = [c for c in X.columns if c not in sig_7_present]
random_33 = np.random.choice(other_proteins, size=min(33, len(other_proteins)), replace=False).tolist()
slice_2_cols = sig_7_present + random_33
slice_2 = final_df[common_keys + slice_2_cols]
slice_2.to_csv('slice_2_sig7_plus_33_random.csv', index=False)

# 3) All proteins
final_df.to_csv('slice_3_all_proteins.csv', index=False)

print("Preprocessing complete. Slices saved.")