app/core/template_generator.py

"""Template generator — joins multiple SAS files, applies column classifications,
and produces paired output files:

  1. Data file: matrix (data variables as rows, samples as columns)
  2. Metadata file: sample-level annotations (one row per sample)
"""

from pathlib import Path

import pandas as pd

from app.core.sas_extractor import read_data
from app.models.mapping import ColumnClassification, FileJoin


def _clean_suffix_label(filename: str) -> str:
    """Strip file extension and optional ::SheetName from a filename for use as merge suffix."""
    # Remove ::SheetName if present
    if "::" in filename:
        base, sheet = filename.rsplit("::", 1)
        # Use sheet name as the label (more descriptive)
        return sheet.replace(" ", "_")
    # Strip any known extension
    for ext in (".sas7bdat", ".csv", ".xpt", ".xlsx", ".xls"):
        if filename.lower().endswith(ext):
            return filename[: -len(ext)]
    return filename


def _coerce_join_keys(left: pd.DataFrame, right: pd.DataFrame, keys: list[str]) -> None:
    """Cast join key columns to the same type in both DataFrames (in-place).

    When one file stores a key as int and another as str, pd.merge fails.
    We coerce both sides to str to guarantee a safe merge.
    """
    for key in keys:
        if key in left.columns and key in right.columns:
            if left[key].dtype != right[key].dtype:
                left[key] = left[key].astype(str).str.strip()
                right[key] = right[key].astype(str).str.strip()


def _load_and_join(
    source_paths: dict[str, Path],
    joins: list[FileJoin],
    primary_file: str,
) -> pd.DataFrame:
    """Load all SAS files and join them based on the AI-determined relationships.

    Args:
        source_paths: mapping of filename -> file path
        joins: how files connect to each other
        primary_file: the main file others join onto

    Returns:
        Single merged DataFrame.
    """
    # Load primary file
    if primary_file not in source_paths:
        raise ValueError(f"Primary file '{primary_file}' not found in uploaded files")

    merged = read_data(source_paths[primary_file])

    # Track which files have been joined
    joined_files = {primary_file}

    # Apply joins in order
    for join in joins:
        # Determine which file to join in
        if join.right_file not in joined_files and join.right_file in source_paths:
            right_df = read_data(source_paths[join.right_file])
            _coerce_join_keys(merged, right_df, join.join_columns)
            merged = pd.merge(
                merged, right_df,
                on=join.join_columns,
                how=join.join_type,
                suffixes=("", f"_{_clean_suffix_label(join.right_file)}"),
            )
            joined_files.add(join.right_file)
        elif join.left_file not in joined_files and join.left_file in source_paths:
            left_df = read_data(source_paths[join.left_file])
            _coerce_join_keys(merged, left_df, join.join_columns)
            merged = pd.merge(
                merged, left_df,
                on=join.join_columns,
                how=join.join_type,
                suffixes=("", f"_{_clean_suffix_label(join.left_file)}"),
            )
            joined_files.add(join.left_file)

    # Load any files that weren't part of any join (standalone files)
    for filename, path in source_paths.items():
        if filename not in joined_files:
            standalone = read_data(path)
            # Try to find common columns for an automatic join
            common = list(set(merged.columns) & set(standalone.columns))
            if common:
                _coerce_join_keys(merged, standalone, common)
                merged = pd.merge(
                    merged, standalone,
                    on=common,
                    how="left",
                    suffixes=("", f"_{_clean_suffix_label(filename)}"),
                )
            # If no common columns, skip (can't join without a key)

    return merged


def _build_paired_output(
    source_paths: dict[str, Path],
    joins: list[FileJoin],
    primary_file: str,
    columns: list[ColumnClassification],
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Join files and split into data matrix + metadata table.

    Supports longitudinal data: when a column is classified as "timepoint",
    the data matrix columns become {sample_id}_{timepoint} to uniquely
    identify each subject×visit combination.

    Returns:
        (data_df, metadata_df)
    """
    df = _load_and_join(source_paths, joins, primary_file)

    # Apply transforms
    for col in columns:
        if col.transform and col.source_column in df.columns:
            df[col.source_column] = _apply_transform(df[col.source_column], col.transform)

    # Identify columns by classification
    sample_id_col = next(c for c in columns if c.classification == "sample_id")
    subject_id_col = next((c for c in columns if c.classification == "subject_id"), None)
    timepoint_col = next((c for c in columns if c.classification == "timepoint"), None)
    metadata_cols = [c for c in columns if c.classification == "metadata"]
    data_cols = [c for c in columns if c.classification == "data"]

    # Output always uses these standard names
    sample_id_name = "Sample_ID"
    subject_id_name = "Subject_title"
    timepoint_name = "Timepoint"

    # Subject ID defaults to Sample ID if not set
    if subject_id_col:
        subject_id_src = subject_id_col.source_column
    else:
        subject_id_src = sample_id_col.source_column  # reuse sample_id

    # --- Build composite column keys for longitudinal data ---
    has_timepoint = timepoint_col and timepoint_col.source_column in df.columns

    if has_timepoint:
        # Create composite key: "{sample_id}_{timepoint}" for each row
        tp_values = df[timepoint_col.source_column].astype(str).str.strip()
        sample_values = df[sample_id_col.source_column].astype(str).str.strip()
        composite_ids = sample_values + "_" + tp_values
    else:
        composite_ids = df[sample_id_col.source_column].astype(str)

    # --- Build metadata file ---
    meta_col_names = [sample_id_col.source_column]
    if subject_id_col and subject_id_col.source_column != sample_id_col.source_column:
        meta_col_names.append(subject_id_col.source_column)
    if has_timepoint:
        meta_col_names.append(timepoint_col.source_column)
    for c in metadata_cols:
        if c.source_column in df.columns:
            meta_col_names.append(c.source_column)
    metadata_df = df[meta_col_names].copy()

    # Rename columns
    rename_meta = {sample_id_col.source_column: sample_id_name}
    if subject_id_col and subject_id_col.source_column != sample_id_col.source_column:
        rename_meta[subject_id_col.source_column] = subject_id_name
    if has_timepoint:
        rename_meta[timepoint_col.source_column] = timepoint_name
    for c in metadata_cols:
        if c.output_name and c.output_name != c.source_column:
            rename_meta[c.source_column] = c.output_name
    metadata_df = metadata_df.rename(columns=rename_meta)

    # Add Subject_title column if not present (copies Sample_ID)
    if subject_id_name not in metadata_df.columns:
        metadata_df.insert(1, subject_id_name, metadata_df[sample_id_name])

    # For longitudinal data, add the composite ID as the first column
    if has_timepoint:
        metadata_df.insert(0, "Composite_ID", composite_ids.values)

    metadata_df = metadata_df.drop_duplicates()

    # --- Build data file ---
    data_src_cols = [c.source_column for c in data_cols if c.source_column in df.columns]
    data_matrix = df[data_src_cols].copy()

    # Rename data columns if output_name specified
    rename_data = {}
    for c in data_cols:
        if c.output_name and c.output_name != c.source_column:
            rename_data[c.source_column] = c.output_name
    data_matrix = data_matrix.rename(columns=rename_data)

    # Handle duplicate column names after merge/rename by appending a numeric suffix
    cols = list(data_matrix.columns)
    seen: dict[str, int] = {}
    for i, col_name in enumerate(cols):
        if col_name in seen:
            seen[col_name] += 1
            cols[i] = f"{col_name}_{seen[col_name]}"
        else:
            seen[col_name] = 0
    data_matrix.columns = cols

    # Transpose: rows become data variables, columns become sample indices
    # For longitudinal data, use composite IDs so each subject×visit is a unique column
    sample_ids = composite_ids.tolist()
    subject_ids = df[subject_id_src].tolist()

    data_matrix.index = sample_ids
    data_df = data_matrix.T

    # The data file format is:
    #   Sample_ID, COVID-01_V1, COVID-01_V2, ...  ← column headers = composite IDs
    #   Subject_title, COVID_M1, COVID_M1, ...     ← first data row with subject titles
    #   Timepoint, Visit 1, Visit 2, ...           ← second data row with timepoints (if longitudinal)
    #   IFIT1, 5.2, 6.1, ...                      ← data variable rows

    data_df.index.name = sample_id_name

    # Insert Subject_title as the first row
    subject_row = pd.DataFrame(
        [subject_ids],
        columns=data_df.columns,
        index=[subject_id_name],
    )
    subject_row.index.name = sample_id_name
    data_df = pd.concat([subject_row, data_df])

    # Insert Timepoint row after Subject_title (if longitudinal)
    if has_timepoint:
        tp_row = pd.DataFrame(
            [df[timepoint_col.source_column].astype(str).tolist()],
            columns=data_df.columns,
            index=[timepoint_name],
        )
        tp_row.index.name = sample_id_name
        # Insert after Subject_title row (position 1)
        data_df = pd.concat([data_df.iloc[:1], tp_row, data_df.iloc[1:]])

    return data_df, metadata_df


def _apply_transform(series: pd.Series, transform: str) -> pd.Series:
    """Apply a simple transform to a pandas Series."""
    t = transform.lower().strip()
    if t in ("uppercase", "upper"):
        return series.astype(str).str.upper()
    elif t in ("lowercase", "lower"):
        return series.astype(str).str.lower()
    elif t in ("strip", "trim"):
        return series.astype(str).str.strip()
    elif t.startswith("date_parse("):
        return pd.to_datetime(series, errors="coerce")
    else:
        return series


def generate_csv(
    source_paths: dict[str, Path],
    joins: list[FileJoin],
    primary_file: str,
    columns: list[ColumnClassification],
    output_dir: Path,
    filename_base: str,
) -> tuple[Path, Path]:
    """Generate paired CSV files: data + metadata."""
    data_df, metadata_df = _build_paired_output(source_paths, joins, primary_file, columns)

    data_path = output_dir / f"data_{filename_base}.csv"
    metadata_path = output_dir / f"metadata_{filename_base}.csv"

    # Column headers = sample IDs, index label = "Sample_ID"
    data_df.to_csv(data_path, index=True)
    metadata_df.to_csv(metadata_path, index=False)

    return data_path, metadata_path


def _load_and_concat(
    source_paths: dict[str, Path],
    included_columns: dict[str, list[str]],
) -> pd.DataFrame:
    """Load all files and join on common columns, keeping only included columns.

    Args:
        source_paths: mapping of filename -> file path
        included_columns: mapping of filename -> list of column names to include

    Returns:
        Single merged DataFrame with only the included columns.
    """
    frames = []
    for filename, path in source_paths.items():
        df = read_data(path)
        # Keep only included columns that exist
        keep = [c for c in included_columns.get(filename, []) if c in df.columns]
        if keep:
            frames.append(df[keep])

    if not frames:
        raise ValueError("No columns selected — nothing to generate")

    if len(frames) == 1:
        return frames[0]

    # Find common columns across all included frames for joining
    col_sets = [set(f.columns) for f in frames]
    common = col_sets[0]
    for cs in col_sets[1:]:
        common = common & cs

    if common:
        # Merge sequentially on common columns
        common_list = list(common)
        merged = frames[0]
        for f in frames[1:]:
            _coerce_join_keys(merged, f, common_list)
            merged = pd.merge(merged, f, on=common_list, how="outer",
                              suffixes=("", "_dup"))
            # Drop duplicate columns from the merge
            merged = merged[[c for c in merged.columns if not c.endswith("_dup")]]
        return merged
    else:
        # No common columns — concatenate side by side
        return pd.concat(frames, axis=1)


def generate_qvey_csv(
    source_paths: dict[str, Path],
    included_columns: dict[str, list[str]],
    output_dir: Path,
    filename_base: str,
) -> Path:
    """Generate a Qvey-format CSV file.

    Output format:
      Row 1: q1, q2, q3, ... (question numbers)
      Row 2: original column headers
      Row 3+: data rows

    Args:
        source_paths: mapping of filename -> file path
        included_columns: mapping of filename -> list of included column names
        output_dir: directory to write the output
        filename_base: base name for the output file

    Returns:
        Path to the generated CSV file.
    """
    df = _load_and_concat(source_paths, included_columns)

    # Row 1 (header): q1, q2, q3, ...
    # Row 2: original column names
    # Row 3+: data
    original_headers = list(df.columns)
    q_headers = [f"q{i+1}" for i in range(len(original_headers))]

    # Rename columns to q1, q2, ... so the CSV header is the question numbers
    df.columns = q_headers

    # Insert original column names as the first data row
    header_row = pd.DataFrame([original_headers], columns=q_headers)
    output = pd.concat([header_row, df], ignore_index=True)

    output_path = output_dir / f"{filename_base}_qvey.csv"
    output.to_csv(output_path, index=False)

    return output_path


def generate_xlsx(
    source_paths: dict[str, Path],
    joins: list[FileJoin],
    primary_file: str,
    columns: list[ColumnClassification],
    output_dir: Path,
    filename_base: str,
) -> Path:
    """Generate Excel workbook with Data + Metadata sheets."""
    data_df, metadata_df = _build_paired_output(source_paths, joins, primary_file, columns)

    output_path = output_dir / f"data_{filename_base}.xlsx"

    with pd.ExcelWriter(output_path, engine="openpyxl") as writer:
        data_df.to_excel(writer, sheet_name="Data", index=True)
        metadata_df.to_excel(writer, sheet_name="Metadata", index=False)

    return output_path