import os
import pandas as pd
import random
import re
from sklearn.preprocessing import MinMaxScaler

# Function to assign main accounts
def assign_main_accounts(creators_file, chatter_files):
    creators = pd.read_excel(creators_file)
    creators.columns = creators.columns.str.strip()

    # Debugging: Check initial columns
    print("DEBUG: Initial Columns in Creator File:", creators.columns)

    # Standardize column names
    column_mapping = {
        "Creator": "Creator",
        "Total earnings": "Total earnings",
        "Subscription": "Subscription",
        "Active Fans": "ActiveFans",
        "Total active fans": "ActiveFans",
    }
    creators.rename(columns={k: v for k, v in column_mapping.items() if k in creators.columns}, inplace=True)

    # Debugging: Check renamed columns
    print("DEBUG: Renamed Columns in Creator File:", creators.columns)

    required_columns = ["Creator", "Total earnings", "Subscription", "ActiveFans"]
    missing_columns = [col for col in required_columns if col not in creators.columns]
    if missing_columns:
        raise KeyError(f"Missing required columns in creators file: {missing_columns}")

    # Process creators file
    creators["Total earnings"] = creators["Total earnings"].replace("[\$,]", "", regex=True).astype(float)
    creators["Subscription"] = creators["Subscription"].replace("[\$,]", "", regex=True).astype(float)
    creators["ActiveFans"] = pd.to_numeric(creators["ActiveFans"], errors="coerce").fillna(0)

    # Normalize data
    scaler = MinMaxScaler()
    creators[["Earnings_Normalized", "Subscriptions_Normalized"]] = scaler.fit_transform(
        creators[["Total earnings", "Subscription"]]
    )
    creators["Penalty Factor"] = 1 - abs(creators["Earnings_Normalized"] - creators["Subscriptions_Normalized"])
    creators["Score"] = (
        0.7 * creators["Earnings_Normalized"] + 0.3 * creators["Subscriptions_Normalized"]
    ) * creators["Penalty Factor"]
    creators["Rank"] = creators["Score"].rank(ascending=False)
    creators = creators.sort_values(by="Rank").reset_index(drop=True)

    processed_creator_file = creators[["Creator", "ActiveFans"]]

    updated_chatter_files = []
    for chatter_file in chatter_files:
        chatters = pd.read_excel(chatter_file)
        chatters.columns = chatters.columns.str.strip()
        if len(chatters) > len(creators):
            raise ValueError("Not enough creators to assign to all chatters.")

        # Assign creators to chatters
        chatters["Main Account"] = creators.iloc[:len(chatters)]["Creator"].values
        updated_chatter_files.append(chatters)

    # Combine all updated chatter files into a single DataFrame
    combined_assignments = pd.concat(updated_chatter_files, ignore_index=True)

    return updated_chatter_files, processed_creator_file, combined_assignments




def save_processed_files(assignments, output_dir):
    """
    Save processed chatter files to the output directory.
    """
    for idx, (shift, data) in enumerate(assignments.items()):
        output_file = os.path.join(output_dir, f"Updated_{shift.lower()}_file.xlsx")
        data.to_excel(output_file, index=False)
        print(f"Saved {shift} file to {output_file}")


# Function to clean chatter data
def clean_chatter_data(chatter_data):
    """
    Clean and prepare chatter data for scheduling.
    """
    required_columns = ["Name", "Main Account", "Final Rating", "Available Work Days"]
    for col in required_columns:
        if col not in chatter_data.columns:
            raise KeyError(f"Missing required column in chatter data: {col}")

    chatter_data["WorkDays"] = pd.to_numeric(chatter_data.get("Available Work Days", 6), errors="coerce").fillna(6).astype(int)
    chatter_data["Desired Off Day"] = chatter_data["Desired Off Day"].fillna("").apply(
        lambda x: [day.strip().capitalize() for day in re.split(r"[ ,]+", x) if day.strip()]
    )

    return chatter_data


# Function to create a blank schedule template
def create_schedule_template(account_data):
    """
    Create a blank schedule template with required columns.
    """
    if "Creator" not in account_data.columns or "ActiveFans" not in account_data.columns:
        raise KeyError("Account data must contain 'Creator' and 'ActiveFans' columns.")

    schedule_template = account_data[["Creator", "ActiveFans"]].copy()
    for day in ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]:
        schedule_template[day] = None

    return schedule_template


# Function to assign main accounts to the schedule
def assign_main_accounts_to_schedule(schedule, chatter_data):
    """
    Assign main accounts to the schedule based on chatter data.
    """
    for _, chatter in chatter_data.iterrows():
        main_account = chatter["Main Account"]
        if main_account in schedule["Creator"].values:
            idx = schedule[schedule["Creator"] == main_account].index[0]
            for day in schedule.columns[2:]:
                schedule.at[idx, day] = chatter["Name"]

    return schedule


# Function to assign off days
def assign_off_days(schedule, chatter_data):
    """
    Assign days off for each chatter based on their 'Desired Off Day' field.
    """
    for _, chatter in chatter_data.iterrows():
        for off_day in chatter["Desired Off Day"]:
            if off_day in schedule.columns[2:]:
                schedule.loc[schedule[off_day] == chatter["Name"], off_day] = None
    return schedule


# Function to randomly fill schedule slots
def randomly_fill_slots(schedule, chatter_data, max_accounts_per_day=3, max_fans_per_day=1000):
    """
    Randomly fill remaining slots in the schedule while respecting constraints.
    """
    days_of_week = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
    chatters_list = chatter_data["Name"].tolist()

    for day in days_of_week:
        for idx, row in schedule.iterrows():
            if pd.isnull(schedule.at[idx, day]):
                random.shuffle(chatters_list)
                for chatter in chatters_list:
                    schedule.at[idx, day] = chatter
                    break

    return schedule


# Main schedule generation function
def generate_schedule(chatter_files, account_data):
    schedules = {}
    shift_names = ["Overnight", "Day", "Prime"]

    for idx, chatter_df in enumerate(chatter_files):
        shift_name = shift_names[idx]
        chatter_df = clean_chatter_data(chatter_df)
        schedule = create_schedule_template(account_data)
        schedule = assign_main_accounts_to_schedule(schedule, chatter_df)
        schedule = assign_off_days(schedule, chatter_df)
        schedule = randomly_fill_slots(schedule, chatter_df)
        schedules[shift_name] = schedule

    return schedules