import os
import numpy as np
import pandas as pd
import gradio as gr
import spaces
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn import metrics

### ------------------------------ ###
###       Data Transformation      ###
### ------------------------------ ###

# Load dataset
uncleaned_data = pd.read_csv('data.csv')

# Remove timestamp from dataset (first column)
uncleaned_data = uncleaned_data.iloc[:, 1:]
data = pd.DataFrame()

# Track categorical mappings
cat_value_dicts = {}
final_colname = uncleaned_data.columns[-1]

for colname, colval in uncleaned_data.items():
    if isinstance(colval.values[0], (np.integer, float, np.floating)):
        data[colname] = uncleaned_data[colname].copy()
        continue

    new_dict = {}
    val = 0
    transformed_col_vals = []

    for item in colval.values:
        if item not in new_dict:
            new_dict[item] = val
            val += 1
        transformed_col_vals.append(new_dict[item])

    if colname == final_colname:
        new_dict = {value: key for key, value in new_dict.items()}

    cat_value_dicts[colname] = new_dict
    data[colname] = transformed_col_vals


### -------------------------------- ###
###           Model Training         ###
### -------------------------------- ###

feature_cols = [c for c in data.columns if c != final_colname]
x = data[feature_cols]
y = data[final_colname]

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25, random_state=42)

# ZeroGPU decorator on startup model training
@spaces.GPU
def train_model(x_tr, y_tr):
    clf = LogisticRegression(max_iter=1000)
    clf.fit(x_tr, y_tr.values.ravel())
    return clf

model = train_model(x_train, y_train)
y_pred = model.predict(x_test)

acc = str(round(metrics.accuracy_score(y_test, y_pred) * 100, 1)) + "%"

def get_feat():
    feats = [abs(val) for val in model.coef_[0]]
    max_val = max(feats)
    idx = feats.index(max_val)
    return feature_cols[idx]

most_imp_feat = get_feat()


### -------------------------------- ###
###        Helper: Read info.md      ###
### -------------------------------- ###

header_text = "Personalized Recommendation AI"
body_text = ""

if os.path.exists('info.md'):
    with open('info.md', 'r') as f:
        lines = f.readlines()
        if lines:
            header_text = lines[0].strip()
            body_text = "".join(lines[1:])


### ------------------------------- ###
###        Interface Creation       ###
### ------------------------------- ###

# ZeroGPU decorator on prediction function
@spaces.GPU
def general_predictor(*args):
    features = []
    for colname, arg in zip(feature_cols, args):
        if colname in cat_value_dicts:
            features.append(cat_value_dicts[colname][arg])
        else:
            features.append(float(arg) if arg is not None else 0.0)

    new_input = [features]
    result = model.predict(new_input)[0]

    if final_colname in cat_value_dicts:
        return cat_value_dicts[final_colname].get(result, str(result))
    return str(result)


with gr.Blocks() as demo:
    gr.Markdown(f"# {header_text}")
    gr.Markdown('Take the quiz to get a personalized recommendation using AI.')

    with gr.Row():
        with gr.Column():
            inputls = []
            for colname in feature_cols:
                if colname in cat_value_dicts:
                    radio_options = list(cat_value_dicts[colname].keys())
                    inputls.append(gr.Dropdown(radio_options, label=colname))
                else:
                    inputls.append(gr.Number(label=colname))

            submit = gr.Button("Click to see your personalized result!", variant="primary")
            output = gr.Textbox(label="Your recommendation:", placeholder="your recommendation will appear here")

            submit.click(fn=general_predictor, inputs=inputls, outputs=output)

            with gr.Row():
                with gr.Column():
                    gr.Markdown(f"### Accuracy:\n{acc}")
                with gr.Column():
                    gr.Markdown(f"### Most important feature:\n{most_imp_feat}")

            gr.Markdown(
                '⭐ *Note that model accuracy is based on the uploaded data.csv and reflects how well the AI model can give correct recommendations for that dataset.*'
            )

        with gr.Column():
            gr.Markdown(body_text)

demo.launch()