"""
COMP 341 Homework 2: Movie Critics

Find critics whose ratings best predict audience scores using data cleaning,
imputation, PCA, and correlation analysis.
"""

import pandas as pd
import numpy as np
from pathlib import Path
from typing import Union, List, Tuple
from sklearn.impute import KNNImputer
from sklearn.decomposition import PCA


# =============================================================================
# Data Loading Functions
# =============================================================================

def load_critic_data(data_dir: Union[str, Path]) -> pd.DataFrame:
    """Load rotten_tomatoes_critic_reviews.csv."""
    data_dir = Path(data_dir)
    return pd.read_csv(data_dir / 'rotten_tomatoes_critic_reviews.csv')


def load_movie_data(data_dir: Union[str, Path]) -> pd.DataFrame:
    """Load rotten_tomatoes_movies.csv."""
    data_dir = Path(data_dir)
    return pd.read_csv(data_dir / 'rotten_tomatoes_movies.csv')


# =============================================================================
# Data Cleaning Functions
# =============================================================================

def remove_unnamed_critics(df: pd.DataFrame) -> pd.DataFrame:
    """Remove critics with no name (empty string or NaN critic_name)."""
    df = df.copy()
    # Remove NaN critic_name
    df = df[df['critic_name'].notna()]
    # Remove empty or whitespace-only critic_name
    df = df[df['critic_name'].str.strip() != '']
    return df.reset_index(drop=True)


def filter_critics_by_review_count(
    df: pd.DataFrame,
    threshold: int = 500
) -> pd.DataFrame:
    """Keep critics with >= threshold scored movies."""
    # Count non-null review_score per critic
    scored_counts = df.groupby('critic_name')['review_score'].apply(
        lambda x: x.notna().sum()
    )
    # Get critics meeting the threshold
    qualifying_critics = scored_counts[scored_counts >= threshold].index
    # Filter to only those critics
    return df[df['critic_name'].isin(qualifying_critics)].reset_index(drop=True)


def _normalize_single_score(score):
    """Convert a single review score to 0-100 scale."""
    # Handle NaN / None
    if score is None or (isinstance(score, float) and np.isnan(score)):
        return np.nan

    score_str = str(score).strip()

    # Handle empty strings
    if score_str == '' or score_str.lower() == 'nan':
        return np.nan

    # Letter grade mapping
    letter_grades = {
        'A+': 95, 'A': 95, 'A-': 92,
        'B+': 88, 'B': 85, 'B-': 82,
        'C+': 78, 'C': 75, 'C-': 72,
        'D+': 68, 'D': 65, 'D-': 62,
        'F+': 55, 'F': 50, 'F-': 45,
        'E': 55,
    }

    # Check for letter grade (case-insensitive)
    upper = score_str.upper().strip()
    if upper in letter_grades:
        return float(letter_grades[upper])

    # Check for fraction
    if '/' in score_str:
        parts = score_str.split('/')
        try:
            numerator = float(parts[0].strip())
            denominator = float(parts[1].strip())
            if denominator == 0:
                return np.nan
            # If numerator > denominator, it's likely already on 0-100 scale
            if numerator > denominator:
                return min(numerator, 100.0)
            result = (numerator / denominator) * 100.0
            return min(result, 100.0)
        except (ValueError, IndexError):
            return np.nan

    # Try to parse as a number
    try:
        val = float(score_str)
        return min(val, 100.0)
    except ValueError:
        return np.nan


def clean_review_scores(df: pd.DataFrame) -> pd.DataFrame:
    """Normalize all review_score values to 0-100 scale."""
    df = df.copy()
    df['review_score'] = df['review_score'].apply(_normalize_single_score)
    return df


# =============================================================================
# Matrix Preparation and Imputation Functions
# =============================================================================

def prepare_critic_movie_matrix(
    critic_df: pd.DataFrame,
    movie_df: pd.DataFrame
) -> pd.DataFrame:
    """Create pivot table: movies (rows) x critics (columns) with audience_score."""
    # Pivot: rows = rotten_tomatoes_link, columns = critic_name, values = review_score
    pivot = critic_df.pivot_table(
        index='rotten_tomatoes_link',
        columns='critic_name',
        values='review_score',
        aggfunc='mean'
    )

    # Join with movie_df to get audience_rating
    movie_scores = movie_df[['rotten_tomatoes_link', 'audience_rating']].copy()
    movie_scores = movie_scores.set_index('rotten_tomatoes_link')

    # Join and rename
    matrix = pivot.join(movie_scores, how='inner')
    matrix = matrix.rename(columns={'audience_rating': 'audience_score'})

    # Drop rows where audience_score is NaN
    matrix = matrix.dropna(subset=['audience_score'])

    return matrix


def impute_missing_zeros(df: pd.DataFrame) -> pd.DataFrame:
    """Fill missing values with 0."""
    return df.fillna(0)


def impute_missing_mean(
    df: pd.DataFrame,
    exclude_cols: List[str] = None
) -> pd.DataFrame:
    """Fill missing values with column mean (per critic)."""
    if exclude_cols is None:
        exclude_cols = ['audience_score']

    result = df.copy()
    for col in result.columns:
        if col not in exclude_cols:
            result[col] = result[col].fillna(result[col].mean())
    return result


def impute_missing_knn(
    df: pd.DataFrame,
    k: int = 5,
    exclude_cols: List[str] = None
) -> pd.DataFrame:
    """Fill missing values using KNN imputation."""
    if exclude_cols is None:
        exclude_cols = ['audience_score']

    result = df.copy()

    # Separate excluded columns
    excluded_data = result[exclude_cols].copy()
    critic_cols = [c for c in result.columns if c not in exclude_cols]
    critic_data = result[critic_cols]

    # Apply KNN imputation
    imputer = KNNImputer(n_neighbors=k)
    imputed_values = imputer.fit_transform(critic_data)

    # Reconstruct DataFrame
    imputed_df = pd.DataFrame(imputed_values, columns=critic_cols, index=result.index)

    # Add back excluded columns
    for col in exclude_cols:
        imputed_df[col] = excluded_data[col]

    return imputed_df


# =============================================================================
# Analysis Functions
# =============================================================================

def calculate_critic_correlation(
    df: pd.DataFrame,
    target_col: str = 'audience_score'
) -> pd.Series:
    """Calculate Pearson correlation between each critic and the target."""
    target = df[target_col]
    critic_cols = [c for c in df.columns if c != target_col]
    correlations = df[critic_cols].corrwith(target)
    # Drop any NaN correlations
    correlations = correlations.dropna()
    return correlations


def get_top_correlated_critics(
    correlations: pd.Series,
    n: int = 5
) -> List[str]:
    """Return the top n critic names most positively correlated with audience score."""
    return correlations.nlargest(n).index.tolist()


def run_pca(
    df: pd.DataFrame,
    n_components: int = 2,
    exclude_cols: List[str] = None
) -> Tuple[np.ndarray, PCA]:
    """Run PCA on the imputed critic-movie matrix."""
    if exclude_cols is None:
        exclude_cols = ['audience_score']

    # Exclude specified columns
    feature_cols = [c for c in df.columns if c not in exclude_cols]
    data = df[feature_cols].values

    pca = PCA(n_components=n_components)
    transformed = pca.fit_transform(data)

    return transformed, pca


# =============================================================================
# Optional: Bias Analysis (Bonus)
# =============================================================================

def calculate_critic_bias(df: pd.DataFrame) -> pd.DataFrame:
    """Calculate fresh/rotten percentages per critic."""
    grouped = df.groupby('critic_name')['review_type'].value_counts().unstack(fill_value=0)

    result = pd.DataFrame()
    result['critic_name'] = grouped.index
    total = grouped.sum(axis=1)
    result['total_reviews'] = total.values

    fresh = grouped.get('Fresh', pd.Series(0, index=grouped.index))
    rotten = grouped.get('Rotten', pd.Series(0, index=grouped.index))

    result['fresh_pct'] = (fresh.values / total.values) * 100
    result['rotten_pct'] = (rotten.values / total.values) * 100

    return result.reset_index(drop=True)