from datetime import datetime


def generate_preprocessing_script(eda: dict) -> str:
    action_plan = eda.get("data_structural_report", {}).get("action_plan", {})
    col_views = eda.get("col_views", {})
    missingness = eda.get("missingness", {})
    categorical_summary = eda.get("categorical_summary", [])
    outliers = eda.get("outliers", {})
    feature_redundancy = eda.get("feature_redundancy", {})

    numeric_cols = col_views.get("numeric_features", [])
    categorical_cols = col_views.get("categorical_features", [])
    identifier_cols = col_views.get("identifiers", [])

    drop_columns = action_plan.get("drop_columns", [])
    log_transform = action_plan.get("apply_log_transform", [])
    log1p_transform = action_plan.get("apply_log1p_transform", [])

    missing_pct = missingness.get("missing_values_percent_cols", {})
    numeric_missing = [c for c in numeric_cols if missing_pct.get(c, 0) > 0]
    categorical_missing = [c for c in categorical_cols if missing_pct.get(c, 0) > 0]

    redundant_pairs = feature_redundancy.get("highly_correlated_pairs", [])
    redundant_to_drop = list({p["feature_1"] for p in redundant_pairs})

    high_cardinality = [
        c["column"] for c in categorical_summary
        if c.get("high_cardinality", False)
    ]
    near_zero_variance = [
        c["column"] for c in categorical_summary
        if c.get("most_frequent_pct", 0) > 0.99
    ]

    lines = []

    # ── Header ──
    lines += [
        "# ============================================================",
        "# Preprocessing Script — Generated by Athena",
        f"# Generated at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "# ============================================================",
        "",
    ]

    # ── Imports ──
    lines += [
        "import pandas as pd",
        "import numpy as np",
        "from sklearn.impute import SimpleImputer",
        "from sklearn.preprocessing import LabelEncoder",
        "",
    ]

    # ── Detected issues as comments ──
    lines += ["# --- What Athena detected ---"]
    if identifier_cols:
        lines.append(f"# Identifier columns (dropped):         {identifier_cols}")
    if drop_columns:
        lines.append(f"# Flagged columns (dropped):             {drop_columns}")
    if redundant_to_drop:
        lines.append(f"# Redundant features (dropped):          {redundant_to_drop}")
    if numeric_missing:
        lines.append(f"# Numeric columns with missing values:   {numeric_missing}")
    if categorical_missing:
        lines.append(f"# Categorical columns with missing vals: {categorical_missing}")
    if log_transform:
        lines.append(f"# Log transform applied to:              {log_transform}")
    if log1p_transform:
        lines.append(f"# Log1p transform applied to:            {log1p_transform}")
    if high_cardinality:
        lines.append(f"# High cardinality columns (review):     {high_cardinality}")
    if near_zero_variance:
        lines.append(f"# Near-zero variance columns (dropped):  {near_zero_variance}")
    lines += ["", ""]

    # ── Main function ──
    lines += [
        "def preprocess(df: pd.DataFrame) -> pd.DataFrame:",
        '    """',
        "    Apply all preprocessing steps detected by Athena.",
        "    Input:  raw DataFrame",
        "    Output: cleaned, transformed DataFrame ready for ML",
        '    """',
        "    df = df.copy()",
        "",
    ]

    # Step 1 — Drop identifiers
    if identifier_cols:
        lines += [
            "    # Step 1: Drop identifier columns",
            f"    identifier_cols = {identifier_cols}",
            "    df.drop(columns=[c for c in identifier_cols if c in df.columns], inplace=True)",
            "",
        ]

    # Step 2 — Drop flagged + redundant columns
    all_drops = list(set(drop_columns + redundant_to_drop))
    if all_drops:
        lines += [
            "    # Step 2: Drop flagged and redundant columns",
            f"    drop_cols = {all_drops}",
            "    df.drop(columns=[c for c in drop_cols if c in df.columns], inplace=True)",
            "",
        ]

    # Step 3 — Drop near zero variance categoricals
    if near_zero_variance:
        lines += [
            "    # Step 3: Drop near-zero variance categorical columns",
            f"    near_zero = {near_zero_variance}",
            "    df.drop(columns=[c for c in near_zero if c in df.columns], inplace=True)",
            "",
        ]

    # Step 4 — Impute numeric missing
    if numeric_missing:
        lines += [
            "    # Step 4: Impute numeric missing values with median",
            f"    numeric_missing = {numeric_missing}",
            "    numeric_missing = [c for c in numeric_missing if c in df.columns]",
            "    if numeric_missing:",
            "        imputer = SimpleImputer(strategy='median')",
            "        df[numeric_missing] = imputer.fit_transform(df[numeric_missing])",
            "",
        ]

    # Step 5 — Impute categorical missing
    if categorical_missing:
        lines += [
            "    # Step 5: Impute categorical missing values with most frequent",
            f"    categorical_missing = {categorical_missing}",
            "    categorical_missing = [c for c in categorical_missing if c in df.columns]",
            "    if categorical_missing:",
            "        cat_imputer = SimpleImputer(strategy='most_frequent')",
            "        df[categorical_missing] = cat_imputer.fit_transform(df[categorical_missing])",
            "",
        ]

    # Step 6 — Log transforms
    if log_transform:
        lines += [
            "    # Step 6: Log transform (use only if all values are strictly positive)",
            f"    log_cols = {log_transform}",
            "    for col in log_cols:",
            "        if col in df.columns:",
            "            if (df[col].dropna() <= 0).any():",
            "                print(f'Warning: {col} has non-positive values, skipping log transform')",
            "            else:",
            "                df[col] = np.log(df[col])",
            "",
        ]

    # Step 7 — Log1p transforms
    if log1p_transform:
        lines += [
            "    # Step 7: Log1p transform (safe for zero values)",
            f"    log1p_cols = {log1p_transform}",
            "    for col in log1p_cols:",
            "        if col in df.columns:",
            "            df[col] = np.log1p(df[col])",
            "",
        ]

    # Step 8 — High cardinality warning + label encode low cardinality
    if categorical_cols:
        safe_encode = [c for c in categorical_cols if c not in high_cardinality and c not in near_zero_variance]
        if safe_encode:
            lines += [
                "    # Step 8: Label encode low-cardinality categorical columns",
                "    # For tree-based models only — use OneHotEncoder for linear models",
                f"    encode_cols = {safe_encode}",
                "    le = LabelEncoder()",
                "    for col in encode_cols:",
                "        if col in df.columns:",
                "            df[col] = df[col].fillna('MISSING').astype(str)",
                "            df[col] = le.fit_transform(df[col])",
                "",
            ]

        if high_cardinality:
            lines += [
                "    # Step 8b: High cardinality columns — manual action required",
                "    # These columns have too many unique values for standard encoding.",
                "    # Options: target encoding, frequency encoding, or drop.",
                f"    # High cardinality columns: {high_cardinality}",
                "    # Example target encoding (requires target y):",
                "    # means = df.groupby(col)[target].mean()",
                "    # df[col] = df[col].map(means)",
                "",
            ]

    # Step 9 — Final cleanup
    lines += [
        "    # Step 9: Final cleanup — reset index",
        "    df.reset_index(drop=True, inplace=True)",
        "",
        "    return df",
        "",
        "",
    ]

    # ── Usage example ──
    lines += [
        "# ── Usage ──────────────────────────────────────────────────",
        "# df_raw = pd.read_csv('your_dataset.csv')",
        "# df_clean = preprocess(df_raw)",
        "# print(df_clean.shape)",
        "# ────────────────────────────────────────────────────────────",
    ]

    return "\n".join(lines)

def generate_nlp_script(eda: dict) -> str:
    from datetime import datetime

    col_views = eda.get("col_views", {})
    free_text = eda.get("free_text", {})
    nlp = eda.get("nlp", {})

    text_cols = [c["column"] for c in free_text.get("free_text_columns", [])]
    identifier_cols = col_views.get("identifiers", [])
    entity_keys = col_views.get("entity_keys", [])
    numeric_cols = col_views.get("numeric_features", [])
    categorical_cols = col_views.get("categorical_features", [])

    vectorization_recs = {}
    if nlp and nlp.get("nlp_analysis"):
        for col_result in nlp["nlp_analysis"]:
            vectorization_recs[col_result["column"]] = col_result["ml_readiness"]["vectorization_recommendation"]

    lines = []

    # ── Header ──
    lines += [
        "# ============================================================",
        "# NLP Preprocessing Script — Generated by Athena",
        f"# Generated at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "# ============================================================",
        "",
    ]

    # ── Imports ──
    lines += [
        "import pandas as pd",
        "import numpy as np",
        "import re",
        "from sklearn.feature_extraction.text import TfidfVectorizer",
        "from sklearn.preprocessing import LabelEncoder",
        "from sklearn.model_selection import train_test_split",
        "import nltk",
        "from nltk.corpus import stopwords",
        "from nltk.stem import WordNetLemmatizer",
        "",
        "# Download required NLTK data (run once)",
        "nltk.download('stopwords', quiet=True)",
        "nltk.download('wordnet', quiet=True)",
        "nltk.download('punkt_tab', quiet=True)",
        "",
    ]

    # ── Detected info as comments ──
    lines += ["# --- What Athena detected ---"]
    if text_cols:
        lines.append(f"# Free text columns:     {text_cols}")
    if identifier_cols:
        lines.append(f"# Identifier columns:    {identifier_cols}")
    if entity_keys:
        lines.append(f"# Entity key columns:    {entity_keys}")
    if numeric_cols:
        lines.append(f"# Numeric features:      {numeric_cols}")
    if categorical_cols:
        lines.append(f"# Categorical features:  {categorical_cols}")
    for col, rec in vectorization_recs.items():
        lines.append(f"# Vectorization ({col}): {rec}")
    lines += ["", ""]

    # ── Cleaning function ──
    lines += [
        "def clean_text(text: str) -> str:",
        '    """Clean a single text string."""',
        "    if not isinstance(text, str):",
        "        return ''",
        "    text = re.sub(r'<[^>]+>', ' ', text)          # Strip HTML tags",
        "    text = re.sub(r'http\\S+|www\\.\\S+', ' ', text)  # Remove URLs",
        "    text = re.sub(r'\\S+@\\S+', ' ', text)           # Remove emails",
        "    text = re.sub(r'[^a-zA-Z0-9\\s]', ' ', text)    # Remove special chars",
        "    text = re.sub(r'\\d+', ' ', text)                # Remove numbers",
        "    text = text.lower().strip()",
        "    text = re.sub(r'\\s+', ' ', text)                # Normalize whitespace",
        "    return text",
        "",
        "",
    ]

    # ── Lemmatization function ──
    lines += [
        "def lemmatize_text(text: str, stop_words: set) -> str:",
        '    """Lemmatize and remove stopwords from text."""',
        "    lemmatizer = WordNetLemmatizer()",
        "    tokens = text.split()",
        "    tokens = [lemmatizer.lemmatize(t) for t in tokens if t not in stop_words and len(t) > 2]",
        "    return ' '.join(tokens)",
        "",
        "",
    ]

    # ── Main preprocess function ──
    lines += [
        "def preprocess(df: pd.DataFrame, target_col: str = None) -> dict:",
        '    """',
        "    Full NLP preprocessing pipeline.",
        "    Input:  raw DataFrame",
        "    Output: dict with 'X' (feature matrix), 'y' (target if provided),",
        "            'vectorizers' (fitted TF-IDF vectorizers per text column)",
        '    """',
        "    df = df.copy()",
        "    stop_words = set(stopwords.words('english'))",
        "",
    ]

    # Step 1 — Drop identifiers and entity keys
    all_drop = list(set(identifier_cols + entity_keys))
    if all_drop:
        lines += [
            "    # Step 1: Drop identifier and entity key columns",
            f"    drop_cols = {all_drop}",
            "    df.drop(columns=[c for c in drop_cols if c in df.columns], inplace=True)",
            "",
        ]

    # Step 2 — Clean text columns
    if text_cols:
        lines += [
            "    # Step 2: Clean text columns",
            f"    text_cols = {text_cols}",
            "    for col in text_cols:",
            "        if col in df.columns:",
            "            df[col] = df[col].fillna('').astype(str)",
            "            df[col] = df[col].apply(clean_text)",
            "            df[col] = df[col].apply(lambda x: lemmatize_text(x, stop_words))",
            "",
        ]

    # Step 3 — Encode categorical columns
    if categorical_cols:
        lines += [
            "    # Step 3: Encode categorical columns",
            f"    categorical_cols = {categorical_cols}",
            "    le = LabelEncoder()",
            "    for col in categorical_cols:",
            "        if col in df.columns and col != target_col:",
            "            df[col] = df[col].fillna('MISSING').astype(str)",
            "            df[col] = le.fit_transform(df[col])",
            "",
        ]

    # Step 4 — TF-IDF vectorization
    if text_cols:
        lines += [
            "    # Step 4: TF-IDF vectorization for text columns",
            "    # Adjust max_features based on your vocab size and available memory",
            "    vectorizers = {}",
            "    tfidf_matrices = []",
            f"    for col in {text_cols}:",
            "        if col in df.columns:",
            "            tfidf = TfidfVectorizer(",
            "                max_features=5000,",
            "                ngram_range=(1, 2),  # unigrams and bigrams",
            "                min_df=2,            # ignore very rare terms",
            "                sublinear_tf=True,   # apply log normalization",
            "            )",
            "            tfidf_matrix = tfidf.fit_transform(df[col])",
            "            vectorizers[col] = tfidf",
            "            tfidf_matrices.append(tfidf_matrix)",
            "            df.drop(columns=[col], inplace=True)  # remove raw text",
            "",
        ]

    # Step 5 — Combine features
    lines += [
        "    # Step 5: Combine TF-IDF features with remaining numeric features",
        "    from scipy.sparse import hstack, csr_matrix",
        "    import scipy.sparse as sp",
        "",
    ]

    if numeric_cols:
        lines += [
            f"    numeric_cols = [c for c in {numeric_cols} if c in df.columns and c != target_col]",
            "    if numeric_cols:",
            "        numeric_matrix = csr_matrix(df[numeric_cols].fillna(0).values)",
            "        tfidf_matrices.append(numeric_matrix)",
            "",
        ]

    lines += [
        "    if tfidf_matrices:",
        "        X = hstack(tfidf_matrices)",
        "    else:",
        "        X = csr_matrix(df.drop(columns=[target_col] if target_col and target_col in df.columns else []).values)",
        "",
    ]

    # Step 6 — Extract target
    lines += [
        "    # Step 6: Extract target",
        "    y = None",
        "    if target_col and target_col in df.columns:",
        "        y = df[target_col].values",
        "",
        "    return {",
        "        'X': X,",
        "        'y': y,",
        "        'vectorizers': vectorizers,",
        "    }",
        "",
        "",
    ]

    # ── Usage example ──
    lines += [
        "# ── Usage ──────────────────────────────────────────────────",
        "# df_raw = pd.read_csv('your_dataset.csv')",
        "# result = preprocess(df_raw, target_col='label')",
        "# X, y = result['X'], result['y']",
        "#",
        "# X_train, X_test, y_train, y_test = train_test_split(",
        "#     X, y, test_size=0.2, random_state=42, stratify=y",
        "# )",
        "#",
        "# # For large vocab, consider upgrading to sentence transformers:",
        "# # from sentence_transformers import SentenceTransformer",
        "# # model = SentenceTransformer('all-MiniLM-L6-v2')",
        "# # embeddings = model.encode(df['text_col'].tolist())",
        "# ────────────────────────────────────────────────────────────",
    ]

    return "\n".join(lines)