"""
COMP 341 Homework 2: Movie Critics

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

Functions to implement:
- load_critic_data: Load critic reviews CSV
- load_movie_data: Load movies CSV
- remove_unnamed_critics: Filter out critics with no name
- filter_critics_by_review_count: Keep critics with >= threshold scored movies
- clean_review_scores: Normalize all scores to 0-100 scale
- prepare_critic_movie_matrix: Create pivot table of critics x movies
- impute_missing_zeros: Fill NaN with 0
- impute_missing_mean: Fill NaN with column mean
- impute_missing_knn: Fill NaN using KNN imputation
- calculate_critic_correlation: Pearson correlation with audience score
- get_top_correlated_critics: Top n critics by correlation
- run_pca: Principal Component 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.

    Args:
        data_dir: Path to directory containing the CSV file

    Returns:
        DataFrame with columns: rotten_tomatoes_link, critic_name, top_critic,
        publisher_name, review_type, review_score, review_date, review_content
    """
    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.

    Args:
        data_dir: Path to directory containing the CSV file

    Returns:
        DataFrame with movie information including audience_rating
    """
    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).

    Args:
        df: Critic reviews DataFrame

    Returns:
        DataFrame with unnamed critics removed
    """
    df = df.copy()
    # Remove NaN critic names
    df = df[df['critic_name'].notna()]
    # Remove empty strings
    df = df[df['critic_name'].str.strip() != '']
    return df


def filter_critics_by_review_count(
    df: pd.DataFrame,
    threshold: int = 500
) -> pd.DataFrame:
    """Filter to keep only critics who have scored at least `threshold` movies.

    Important: Count movies with actual scores (non-NaN review_score),
    not just total reviews.

    Args:
        df: Critic reviews DataFrame
        threshold: Minimum number of scored movies

    Returns:
        DataFrame containing only reviews from critics who meet the threshold
    """
    # Count non-null review_score per critic
    scored_counts = df[df['review_score'].notna()].groupby('critic_name').size()
    qualified_critics = scored_counts[scored_counts >= threshold].index
    return df[df['critic_name'].isin(qualified_critics)]


def clean_review_scores(df: pd.DataFrame) -> pd.DataFrame:
    """Normalize all review_score values to 0-100 scale.

    Must handle these formats:
    - Fractions: "3/5" -> 60, "3/4" -> 75, "87/100" -> 87, "4/10" -> 40
    - Letter grades: "A" -> 95, "A-" -> 92, "B+" -> 88, "B" -> 85,
                     "B-" -> 82, "C+" -> 78, "C" -> 75, "C-" -> 72,
                     "D+" -> 68, "D" -> 65, "D-" -> 62, "F" -> 50
    - Direct numbers: Already in 0-100 range, keep as-is
    - Empty/NaN: Preserve as NaN (do not convert to 0)

    Edge cases to handle:
    - Fractions where numerator > denominator (e.g., "92/10"):
      These are likely already on a 0-100 scale, so use numerator directly
      (92/10 -> 92, not 9.2*100=920)
    - Values > 100: cap at 100 (likely data entry errors)

    Args:
        df: Critic reviews DataFrame with review_score column

    Returns:
        DataFrame with normalized review_score column (float, 0-100 or NaN)
    """
    LETTER_GRADES = {
        'A': 95, 'A-': 92,
        'B+': 88, 'B': 85, 'B-': 82,
        'C+': 78, 'C': 75, 'C-': 72,
        'D+': 68, 'D': 65, 'D-': 62,
        'F': 50
    }

    def convert_score(score):
        if pd.isna(score):
            return np.nan
        score = str(score).strip()
        if score == '' or score.lower() == 'nan':
            return np.nan

        # Letter grades (check before numeric to avoid confusion)
        if score in LETTER_GRADES:
            return float(LETTER_GRADES[score])

        # Fractions
        if '/' in score:
            parts = score.split('/')
            if len(parts) == 2:
                try:
                    num = float(parts[0].strip())
                    den = float(parts[1].strip())
                    if den == 0:
                        return np.nan
                    # Edge case: numerator > denominator -> use numerator directly
                    if num > den:
                        result = num
                    else:
                        result = (num / den) * 100
                    return min(result, 100.0)
                except (ValueError, ZeroDivisionError):
                    return np.nan

        # Direct numbers
        try:
            val = float(score)
            return min(val, 100.0)
        except ValueError:
            return np.nan

    df = df.copy()
    df['review_score'] = df['review_score'].apply(convert_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.

    Args:
        critic_df: Cleaned critic reviews DataFrame (with normalized scores)
        movie_df: Movies DataFrame with audience_rating

    Returns:
        DataFrame where:
        - Each row is a movie (indexed by rotten_tomatoes_link)
        - Each column is a critic (their normalized score for that movie)
        - Last column is 'audience_score' from movie_df
        - Missing values are NaN (critic didn't review that movie)
    """
    # Create pivot table: rows=movies, columns=critics, values=scores
    pivot = critic_df.pivot_table(
        index='rotten_tomatoes_link',
        columns='critic_name',
        values='review_score',
        aggfunc='mean'
    )
    pivot.columns.name = None  # Remove column name

    # Get audience ratings from movie_df
    movie_subset = movie_df[['rotten_tomatoes_link', 'audience_rating']].copy()
    movie_subset = movie_subset.set_index('rotten_tomatoes_link')
    movie_subset = movie_subset.rename(columns={'audience_rating': 'audience_score'})

    # Join on rotten_tomatoes_link
    result = pivot.join(movie_subset, how='inner')

    return result


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

    Args:
        df: Critic-movie matrix with NaN values

    Returns:
        DataFrame with NaN replaced by 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).

    Args:
        df: Critic-movie matrix with NaN values
        exclude_cols: Columns to exclude from imputation (e.g., ['audience_score'])

    Returns:
        DataFrame with NaN replaced by column means
    """
    if exclude_cols is None:
        exclude_cols = ['audience_score']

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


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

    Uses k nearest neighbors to estimate missing values based on
    similar movies' scores.

    Args:
        df: Critic-movie matrix with NaN values
        k: Number of neighbors for KNN

    Returns:
        DataFrame with NaN replaced by KNN-imputed values
    """
    if exclude_cols is None:
        exclude_cols = ['audience_score']

    # Separate excluded columns
    excluded_data = {col: df[col].copy() for col in exclude_cols if col in df.columns}
    critic_cols = [c for c in df.columns if c not in exclude_cols]

    # Apply KNN imputation to critic columns only
    imputer = KNNImputer(n_neighbors=k)
    imputed_values = imputer.fit_transform(df[critic_cols])

    result = pd.DataFrame(imputed_values, index=df.index, columns=critic_cols)

    # Add back excluded columns
    for col, data in excluded_data.items():
        result[col] = data

    return result


# =============================================================================
# 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.

    Args:
        df: Imputed critic-movie matrix (no NaN values)
        target_col: Column to correlate against (default: audience_score)

    Returns:
        Series with critic names as index and correlation coefficients as values
    """
    target = df[target_col]
    critic_cols = [c for c in df.columns if c != target_col]
    correlations = df[critic_cols].corrwith(target)
    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.

    Args:
        correlations: Series of correlation values (critic name -> correlation)
        n: Number of top critics to return

    Returns:
        List of critic names, sorted by correlation (highest first)
    """
    return list(correlations.nlargest(n).index)


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.

    Args:
        df: Imputed critic-movie matrix
        n_components: Number of principal components
        exclude_cols: Columns to exclude (e.g., ['audience_score'])

    Returns:
        Tuple of:
        - Transformed data (n_samples, n_components)
        - Fitted PCA object (for explained_variance_ratio_)
    """
    if exclude_cols is None:
        exclude_cols = ['audience_score']

    critic_cols = [c for c in df.columns if c not in exclude_cols]
    pca = PCA(n_components=n_components)
    transformed = pca.fit_transform(df[critic_cols])
    return transformed, pca


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

def calculate_critic_bias(df: pd.DataFrame) -> pd.DataFrame:
    """Calculate fresh/rotten percentages per critic.

    Args:
        df: Critic reviews DataFrame with review_type column

    Returns:
        DataFrame with columns: critic_name, fresh_pct, rotten_pct, total_reviews
    """
    bias = df.groupby('critic_name')['review_type'].value_counts().unstack(fill_value=0)
    bias.columns.name = None

    # Ensure Fresh and Rotten columns exist
    if 'Fresh' not in bias.columns:
        bias['Fresh'] = 0
    if 'Rotten' not in bias.columns:
        bias['Rotten'] = 0

    bias['total_reviews'] = bias.sum(axis=1)
    bias['fresh_pct'] = bias['Fresh'] / bias['total_reviews'] * 100
    bias['rotten_pct'] = bias['Rotten'] / bias['total_reviews'] * 100

    result = bias[['fresh_pct', 'rotten_pct', 'total_reviews']].reset_index()
    return result