import gradio as gr
import pandas as pd
import numpy as np
import tempfile
import os


def calculate_nps(scores):
    scores = np.array(scores)
    total = len(scores)
    if total == 0:
        return 0
    promoters = np.sum((scores == 4) | (scores == 5))
    detractors = np.sum((scores == 1) | (scores == 2))
    return (promoters - detractors) / total * 100


def adjust_nps(file, last_data_row, desired_nps_row, progress=gr.Progress(track_tqdm=False)):
    if file is None:
        raise gr.Error("Please upload an Excel file.")

    if last_data_row is None or desired_nps_row is None:
        raise gr.Error("Please enter both row numbers.")

    try:
        progress(0, desc="Reading file...")
        df = pd.read_excel(file.name, sheet_name="Data", header=None)
    except Exception as e:
        raise gr.Error(f"Could not read sheet named 'Data'. Error: {e}")

    # Convert user-friendly row numbers (e.g., 1503) to 0-based indices
    individual_score_start = 3          # as in your original script
    individual_score_end = int(last_data_row) - 1
    desired_nps_row_idx = int(desired_nps_row) - 1

    if individual_score_end <= individual_score_start:
        raise gr.Error("LAST data record row must be greater than 4 (since data starts at row 4).")

    last_column_index = df.shape[1] - 1
    total_columns = last_column_index + 1

    for col in range(0, last_column_index + 1):
        progress(col / total_columns, desc=f"Processing column {col + 1}/{total_columns}")

        # Get desired NPS for this column
        try:
            desired_nps = float(df.at[desired_nps_row_idx, col])
        except (ValueError, TypeError):
            continue  # skip if not a valid desired NPS
        except KeyError:
            continue

        original_scores = df.loc[individual_score_start:individual_score_end, col].copy()
        scores = pd.to_numeric(original_scores, errors="coerce").dropna()

        if len(scores) < 10:
            # Not enough data, skip column
            continue

        neutral_cap = np.random.uniform(3, 12)
        current_nps = calculate_nps(scores)
        changes = 0
        max_changes = 400

        while abs(current_nps - desired_nps) > 0.1 and changes < max_changes:
            neutral_percent = (scores == 3).sum() / len(scores) * 100

            if current_nps > desired_nps:
                # Reduce NPS
                candidates = scores[(scores == 5) | (scores == 4)].index
                if len(candidates):
                    idx = np.random.choice(candidates)
                    scores.loc[idx] = 4 if scores.loc[idx] == 5 else 3
                    changes += 1
                elif neutral_percent > 0:
                    candidates = scores[scores == 3].index
                    if len(candidates):
                        idx = np.random.choice(candidates)
                        scores.loc[idx] = 2
                        changes += 1
                    else:
                        break
                else:
                    break
            else:
                # Increase NPS
                if neutral_percent < neutral_cap:
                    candidates = scores[scores == 2].index
                    if len(candidates):
                        idx = np.random.choice(candidates)
                        scores.loc[idx] = 3
                        changes += 1
                    else:
                        candidates = scores[scores == 3].index
                        if len(candidates):
                            idx = np.random.choice(candidates)
                            scores.loc[idx] = 4
                            changes += 1
                        else:
                            break
                else:
                    candidates = scores[scores == 3].index
                    if len(candidates):
                        idx = np.random.choice(candidates)
                        scores.loc[idx] = 4
                        changes += 1
                    else:
                        break

            current_nps = calculate_nps(scores)

        # Write back adjusted scores
        df.loc[individual_score_start:individual_score_end, col] = scores

    progress(1, desc="Writing adjusted file...")

    # Save output to a temp file and return path
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, "NeuroSinQ_NPS_Adjusted.xlsx")
    df.to_excel(output_path, index=False, header=False)

    return output_path


with gr.Blocks(title="NeuroSinQ NPS Adjuster") as demo:
    gr.Markdown(
        """
        # 🧠 NeuroSinQ NPS Adjuster

        1. Upload an Excel file with **one worksheet** named **`Data`**.  
        2. Enter the **row number of LAST data record**.  
        3. Enter the **row number of Desired NPS**.  
        4. Click **Make Adjustment** to download the adjusted file.
        """
    )

    with gr.Row():
        file_input = gr.File(
            label="Upload Excel File (.xlsx)",
            file_types=[".xlsx"],
        )

    with gr.Row():
        last_row = gr.Number(
            label="Row number of LAST data record (e.g., 1503)",
            precision=0
        )
        desired_row = gr.Number(
            label="Row number for Desired NPS (e.g., 1523)",
            precision=0
        )

    run_btn = gr.Button("⚙️ Make Adjustment")

    output_file = gr.File(
        label="Download Adjusted File",
        interactive=False
    )

    run_btn.click(
        fn=adjust_nps,
        inputs=[file_input, last_row, desired_row],
        outputs=output_file
    )


if __name__ == "__main__":
    demo.launch(ssr_mode=False)