#!/usr/bin/env python3

################################################################################
# SCRIPT 06: REPEAT PROFILING - Profile Repeat Composition from RepeatMasker
################################################################################

# PURPOSE:
# Parses RepeatMasker .out files from script 05 to determine repeat
# composition fingerprint for each sample.

# INPUT:
# - RepeatMasker .out files from outputs/05_repeatmasker/
# - Can also accept pre-computed .out files via --input-repeatmasker-out

# OUTPUT:
# - repeat_summary_pivoted.tsv (sample × repeat_class/family composition)
# - repeat_summary_long.tsv (detailed long-format data)
# - repeat_statistics.json (comprehensive statistics per sample)
# - repeat_annotations_raw.tsv (complete parsed annotations)

# USAGE:
# python 06_repeat_profiling.py [options]

# OPTIONS (auto-detection from Phase 1):
# python 06_repeat_profiling.py
#   Auto-detects .out files from outputs/05_repeatmasker/

# OPTIONS (explicit specification):
# python 06_repeat_profiling.py \
#   --input-repeatmasker-out RCA outputs/05_repeatmasker/rca.fasta.out \
#   --input-repeatmasker-out gDNA-NT outputs/05_repeatmasker/gdna-nt.fasta.out \
#   --output-dir outputs/06_repeat_profiling

# EXAMPLES:
# # Auto-detect from Phase 1 outputs
# python 06_repeat_profiling.py

# # Explicit input files
# python 06_repeat_profiling.py \
#   --input-repeatmasker-out RCA outputs/05_repeatmasker/rca.fasta.out \
#   --input-repeatmasker-out gDNA-NT outputs/05_repeatmasker/gdna-nt.fasta.out \
#   --output-dir outputs/06_repeat_profiling

# DEPENDENCIES:
# - pandas
# - numpy

# ALGORITHM (EXACT MATCH TO ORIGINAL):
# 1. Locate RepeatMasker .out files (auto-detect or explicit)
# 2. For each .out file:
#    a. Parse with robust error handling (original algorithm)
#    b. Extract class/family with split('/', 1) preservation
#    c. Calculate total masked base pairs
#    d. Calculate percentage per class/family
# 3. Aggregate across all samples
# 4. Output in multiple formats:
#    - Pivoted wide format (for visualization tools)
#    - Long format (for statistical analysis)
#    - JSON statistics (for programmatic access)
#    - Raw annotations (for detailed inspection)

################################################################################

import os
import sys
import argparse
import logging
import glob
import json
from pathlib import Path
from collections import defaultdict

import pandas as pd
import numpy as np

# ============================================================================
# LOGGING SETUP
# ============================================================================

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ============================================================================
# REPEATMASKER PARSING
# ============================================================================

def parse_repeatmasker_output(rm_out_file):
    """
    Parses RepeatMasker .out file
    
    Args:
        rm_out_file (str): Path to RepeatMasker .out file
        
    Returns:
        pd.DataFrame: Parsed annotations with complete schema
    """
    annotations = []
    line_count = 0
    parsed_count = 0
    skipped_count = 0
    
    logger.info(f"Parsing RepeatMasker output: {rm_out_file}")
    
    try:
        with open(rm_out_file, 'r') as f:
            for line_num, line in enumerate(f):
                line_count += 1
                line = line.strip()
                
                # Skip header lines
                if not line or line.startswith("SW") or line.startswith("score") or \
                   line.startswith("There") or line_num < 3:
                    skipped_count += 1
                    continue
                
                parts = line.split()
                
                # Check minimum required columns
                if len(parts) < 11:
                    logger.warning(
                        f"Skipping malformed line #{line_num+1} in {rm_out_file} "
                        f"(too few columns): {line}"
                    )
                    skipped_count += 1
                    continue
                
                try:
                    # Standard RepeatMasker .out columns
                    # 0: SW score
                    # 4: Query sequence name (read_id)
                    # 5: Query start
                    # 6: Query end
                    # 8: Strand (+/C)
                    # 9: Matching repeat name (e.g., L1HS, AluY, MER41B, ALR, (GA)n)
                    # 10: Repeat class/family (e.g., LINE/L1, SINE/Alu, DNA/hAT-Tip100)
                    
                    score = int(parts[0])
                    read_id = parts[4]
                    start = int(parts[5])
                    end = int(parts[6])
                    strand = parts[8]
                    repeat_name = parts[9]
                    repeat_class_family_raw = parts[10]
                    
                    # Determine Class
                    # Use split('/', 1) to preserve complex names
                    if '/' in repeat_class_family_raw:
                        repeat_class = repeat_class_family_raw.split('/', 1)[0]
                    else:
                        # If no '/', the whole string is the class
                        repeat_class = repeat_class_family_raw
                    
                    # Determine Family
                    if '/' in repeat_class_family_raw:
                        # If 'class/family' format, use the part after '/'
                        repeat_family = repeat_class_family_raw.split('/', 1)[1]
                    else:
                        # If only 'class' is given (no '/'), use the specific
                        # repeat_name as the family identifier
                        # This distinguishes between different Simple_repeats
                        # (e.g., (GA)n, (CAT)n) or different Satellites
                        repeat_family = repeat_name
                    
                    # Strand conversion
                    strand_simple = '+' if strand == '+' else '-'
                    
                    # Calculate hit length
                    length = abs(end - start) + 1
                    
                    # Complete data structure
                    data = {
                        "read_id": read_id,
                        "start": start,
                        "end": end,
                        "strand": strand_simple,
                        "repeat_name": repeat_name,
                        "repeat_class_family_raw": repeat_class_family_raw,
                        "repeat_class": repeat_class,
                        "repeat_family": repeat_family,
                        "score": score,
                        "length": length,
                        "source_tool": "RepeatMasker"
                    }
                    
                    annotations.append(data)
                    parsed_count += 1
                    
                except (ValueError, IndexError) as e:
                    logger.warning(
                        f"Skipping line #{line_num+1} due to parsing error ({e}) "
                        f"in {rm_out_file}: {line}"
                    )
                    skipped_count += 1
                    
    except FileNotFoundError:
        logger.error(f"File not found: {rm_out_file}")
        return pd.DataFrame()
    except Exception as e:
        logger.error(f"Error parsing RepeatMasker file {rm_out_file}: {e}")
        raise
    
    # Detailed logging
    logger.info(
        f"Finished parsing {rm_out_file}. Total lines: {line_count}, "
        f"Parsed annotations: {parsed_count}, Skipped/Header lines: {skipped_count}"
    )
    
    return pd.DataFrame(annotations) if annotations else pd.DataFrame()

# ============================================================================
# REPEAT COMPOSITION CALCULATION
# ============================================================================

def calculate_repeat_composition(annotations_df, sample_name):
    """
    Calculate repeat class and family composition percentages.
    
    percentages calculated relative to
    total masked base pairs (not total genome size).
    
    Args:
        annotations_df (pd.DataFrame): Parsed annotations
        sample_name (str): Sample identifier
        
    Returns:
        dict: Complete statistics including class, family, and total BP
    """
    if annotations_df.empty:
        logger.warning(f"No annotations for sample {sample_name}")
        return {
            'sample_name': sample_name,
            'class': {},
            'family': {},
            'total_bp_masked': 0,
            'num_annotations': 0
        }
    
    # Ensure length is numeric
    annotations_df['length'] = pd.to_numeric(
        annotations_df['length'], errors='coerce'
    ).fillna(0).astype(int)
    
    # Calculate total masked base pairs 
    total_masked_bp = annotations_df['length'].sum()
    
    logger.info(
        f"Sample {sample_name}: Total masked BP = {total_masked_bp:,}, "
        f"Total annotations = {len(annotations_df):,}"
    )
    
    if total_masked_bp == 0:
        logger.warning(f"No base pairs masked for sample {sample_name}")
        return {
            'sample_name': sample_name,
            'class': {},
            'family': {},
            'total_bp_masked': 0,
            'num_annotations': len(annotations_df)
        }
    
    # Aggregate by class
    class_stats = {}
    for repeat_class in annotations_df['repeat_class'].unique():
        class_df = annotations_df[annotations_df['repeat_class'] == repeat_class]
        bp_covered = class_df['length'].sum()
        
        # Percentage relative to total masked BP
        percentage = (bp_covered / total_masked_bp) * 100
        
        class_stats[repeat_class] = {
            'bp_covered': int(bp_covered),
            'percentage': float(percentage),
            'num_hits': len(class_df)
        }
    
    # Aggregate by family
    family_stats = {}
    for repeat_family in annotations_df['repeat_family'].unique():
        family_df = annotations_df[annotations_df['repeat_family'] == repeat_family]
        bp_covered = family_df['length'].sum()
        
        # Percentage relative to total masked BP
        percentage = (bp_covered / total_masked_bp) * 100
        
        family_stats[repeat_family] = {
            'bp_covered': int(bp_covered),
            'percentage': float(percentage),
            'num_hits': len(family_df)
        }
    
    return {
        'sample_name': sample_name,
        'class': class_stats,
        'family': family_stats,
        'total_bp_masked': int(total_masked_bp),
        'num_annotations': len(annotations_df)
    }

# ============================================================================
# AUTO-DETECTION OF INPUT FILES
# ============================================================================

def auto_detect_repeatmasker_files(input_dir="outputs/05_repeatmasker"):
    """
    Auto-detect RepeatMasker .out files from Phase 1 outputs.
    
    Searches for *.out and *.fasta.out files in the specified directory.
    
    Args:
        input_dir (str): Directory to search for .out files
        
    Returns:
        dict: Mapping of sample_name -> file_path
    """
    samples = {}
    
    if not os.path.exists(input_dir):
        logger.warning(f"Input directory not found: {input_dir}")
        return samples
    
    logger.info(f"Auto-detecting .out files in: {input_dir}")
    
    # Search for both *.out and *.fasta.out patterns
    patterns = [
        os.path.join(input_dir, "*.out"),
        os.path.join(input_dir, "*.fasta.out")
    ]
    
    for pattern in patterns:
        for out_file in glob.glob(pattern):
            basename = os.path.basename(out_file)
            
            # Extract sample name (e.g., "rca.fasta.out" -> "rca")
            sample_name = basename.replace('.fasta.out', '').replace('.out', '')
            
            # Avoid duplicates
            if sample_name not in samples:
                samples[sample_name] = out_file
                logger.info(f"  Found: {sample_name} <- {out_file}")
    
    return samples

# ============================================================================
# OUTPUT GENERATION
# ============================================================================

def create_long_format_data(all_compositions):
    """
    Create long-format DataFrame for statistical analysis.
    
    Args:
        all_compositions (dict): Complete composition data per sample
        
    Returns:
        pd.DataFrame: Long-format data
    """
    long_data = []
    
    for sample_name, composition in all_compositions.items():
        # Add class entries
        for repeat_class, stats in composition['class'].items():
            long_data.append({
                'sample_name': sample_name,
                'category_type': 'Class',
                'category_name': repeat_class,
                'bp_covered': stats['bp_covered'],
                'percentage_of_total_bp': stats['percentage'],
                'num_hits': stats['num_hits']
            })
        
        # Add family entries
        for repeat_family, stats in composition['family'].items():
            long_data.append({
                'sample_name': sample_name,
                'category_type': 'Family',
                'category_name': repeat_family,
                'bp_covered': stats['bp_covered'],
                'percentage_of_total_bp': stats['percentage'],
                'num_hits': stats['num_hits']
            })
    
    return pd.DataFrame(long_data)

def create_pivoted_summary(long_df):
    """
    Create pivoted wide-format summary output.
    
    Args:
        long_df (pd.DataFrame): Long-format data
        
    Returns:
        pd.DataFrame: Pivoted wide-format data
    """
    if long_df.empty:
        logger.warning("No data to pivot")
        return pd.DataFrame()
    
    # Pivot with multi-level columns
    pivot_df = long_df.pivot_table(
        index='sample_name',
        columns=['category_type', 'category_name'],
        values='percentage_of_total_bp',
        fill_value=0.0
    )
    
    # Flatten multi-index columns: ('Class', 'LINE') -> 'Class_LINE'
    pivot_df.columns = ['_'.join(col).strip() for col in pivot_df.columns.values]
    pivot_df.reset_index(inplace=True)
    
    return pivot_df

def create_statistics_json(all_compositions):
    """
    Create JSON statistics file
    
    Args:
        all_compositions (dict): Complete composition data per sample
        
    Returns:
        dict: Statistics suitable for JSON serialization
    """
    stats_data = {}
    
    for sample_name, composition in all_compositions.items():
        stats_data[sample_name] = {
            'total_bp_masked': composition['total_bp_masked'],
            'num_annotations': composition['num_annotations'],
            'num_classes': len(composition['class']),
            'num_families': len(composition['family']),
            'classes': composition['class'],
            'families': composition['family']
        }
    
    return stats_data

# ============================================================================
# JSON ENCODER FOR NUMPY TYPES
# ============================================================================

class NpEncoder(json.JSONEncoder):
    """JSON encoder that handles numpy types."""
    def default(self, obj):
        if isinstance(obj, (np.integer, np.int64, np.int32, np.int16, np.int8)):
            return int(obj)
        if isinstance(obj, (np.floating, np.float64, np.float32)):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NpEncoder, self).default(obj)

# ============================================================================
# MAIN ANALYSIS
# ============================================================================

def main():
    parser = argparse.ArgumentParser(
        description=(
            "Profile repeat composition from RepeatMasker output. "
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
EXAMPLES:
  # Auto-detect .out files from script 05
  python 06_repeat_profiling.py
  
  # Explicit .out files
  python 06_repeat_profiling.py \\
    --input-repeatmasker-out RCA outputs/05_repeatmasker/rca.fasta.out \\
    --input-repeatmasker-out gDNA-NT outputs/05_repeatmasker/gdna-nt.fasta.out
  
  # Custom output directory
  python 06_repeat_profiling.py \\
    --input-dir outputs/05_repeatmasker \\
    --output-dir outputs/06_repeat_profiling

OUTPUT FILES:
  - repeat_summary_pivoted.tsv     Wide format for visualization
  - repeat_summary_long.tsv         Long format for statistics
  - repeat_statistics.json          Complete per-sample statistics
  - repeat_annotations_raw.tsv      All parsed annotations

NEXT STEP:
  python scripts/07_repeat_analysis.py
        """
    )
    
    # Input specification
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument(
        "--input-dir",
        default="outputs/05_repeatmasker",
        help="Directory containing .out files (auto-detection mode). Default: %(default)s"
    )
    
    input_group.add_argument(
        "--input-repeatmasker-out",
        nargs=2,
        action='append',
        metavar=('SAMPLE_NAME', 'OUT_FILE'),
        help=(
            "Explicitly specify .out files. Can be used multiple times. "
            "Example: --input-repeatmasker-out RCA /path/to/rca.out "
            "--input-repeatmasker-out gDNA /path/to/gdna.out"
        )
    )
    
    # Output
    parser.add_argument(
        "--output-dir",
        default="outputs/06_repeat_profiling",
        help="Output directory. Default: %(default)s"
    )
    
    parser.add_argument(
        "--output-summary",
        default="repeat_summary_pivoted.tsv",
        help="Output summary filename (in --output-dir). Default: %(default)s"
    )
    
    parser.add_argument(
        "--save-raw-annotations",
        action='store_true',
        help="Save complete parsed annotations to TSV (large file)"
    )
    
    args = parser.parse_args()
    
    logger.info("="*70)
    logger.info("REPEAT PROFILING (Script 06)")
    logger.info("="*70)
    logger.info("")
    
    # ========================================================================
    # DETERMINE INPUT FILES
    # ========================================================================
    
    samples_dict = {}
    
    if args.input_repeatmasker_out:
        # Explicit specification
        logger.info("Using explicitly specified .out files:")
        for sample_name, file_path in args.input_repeatmasker_out:
            if not os.path.exists(file_path):
                logger.error(f"File not found: {file_path}")
                sys.exit(1)
            samples_dict[sample_name] = file_path
            logger.info(f"  {sample_name}: {file_path}")
    else:
        # Auto-detection
        samples_dict = auto_detect_repeatmasker_files(args.input_dir)
        
        if not samples_dict:
            logger.error(f"No .out files found in {args.input_dir}")
            logger.info("Use --input-repeatmasker-out to specify files explicitly")
            sys.exit(1)
    
    logger.info(f"\nProcessing {len(samples_dict)} sample(s)")
    logger.info("")
    
    # ========================================================================
    # CREATE OUTPUT DIRECTORY
    # ========================================================================
    
    os.makedirs(args.output_dir, exist_ok=True)
    logger.info(f"Output directory: {args.output_dir}")
    logger.info("")
    
    # ========================================================================
    # PARSE ALL SAMPLES
    # ========================================================================
    
    all_compositions = {}
    all_annotations = []
    
    for sample_name, out_file in samples_dict.items():
        logger.info("="*70)
        logger.info(f"Processing sample: {sample_name}")
        logger.info("="*70)
        
        annotations_df = parse_repeatmasker_output(out_file)
        
        if annotations_df.empty:
            logger.warning(f"No annotations found for {sample_name}")
            all_compositions[sample_name] = {
                'sample_name': sample_name,
                'class': {},
                'family': {},
                'total_bp_masked': 0,
                'num_annotations': 0
            }
            continue
        
        # Add sample_name column for aggregation
        annotations_df['sample_name'] = sample_name
        all_annotations.append(annotations_df)
        
        # Calculate composition
        composition = calculate_repeat_composition(annotations_df, sample_name)
        all_compositions[sample_name] = composition
        
        # Log summary
        logger.info(f"  Total BP masked: {composition['total_bp_masked']:,}")
        logger.info(f"  Total annotations: {composition['num_annotations']:,}")
        logger.info(f"  Unique repeat classes: {len(composition['class'])}")
        logger.info(f"  Unique repeat families: {len(composition['family'])}")
        
        # Log top classes
        if composition['class']:
            top_classes = sorted(
                composition['class'].items(),
                key=lambda x: x[1]['percentage'],
                reverse=True
            )[:5]
            
            logger.info("  Top 5 repeat classes:")
            for cls, stats in top_classes:
                logger.info(
                    f"    {cls}: {stats['percentage']:.2f}% "
                    f"({stats['bp_covered']:,} bp, {stats['num_hits']:,} hits)"
                )
        
        logger.info("")
    
    # ========================================================================
    # CREATE OUTPUT FILES
    # ========================================================================
    
    logger.info("="*70)
    logger.info("Creating output files")
    logger.info("="*70)
    
    # 1. Long-format summary
    long_df = create_long_format_data(all_compositions)
    long_file = os.path.join(args.output_dir, "repeat_summary_long.tsv")
    
    if not long_df.empty:
        long_df.to_csv(long_file, sep='\t', index=False, float_format='%.6f')
        logger.info(f"✓ Long-format summary: {long_file}")
        logger.info(f"  Rows: {len(long_df):,}")
    else:
        logger.warning("No data for long-format summary")
    
    # 2. Pivoted wide-format summary
    pivot_file = os.path.join(args.output_dir, args.output_summary)
    
    if not long_df.empty:
        pivot_df = create_pivoted_summary(long_df)
        pivot_df.to_csv(pivot_file, sep='\t', index=False, float_format='%.6f')
        logger.info(f"✓ Pivoted summary: {pivot_file}")
        logger.info(f"  Samples: {len(pivot_df)}, Columns: {len(pivot_df.columns)}")
    else:
        logger.warning("No data for pivoted summary")
    
    # 3. JSON statistics
    stats_file = os.path.join(args.output_dir, "repeat_statistics.json")
    stats_data = create_statistics_json(all_compositions)
    
    with open(stats_file, 'w') as f:
        json.dump(stats_data, f, indent=2, cls=NpEncoder)
    logger.info(f"✓ Statistics JSON: {stats_file}")
    
    # 4. Raw annotations
    if args.save_raw_annotations and all_annotations:
        raw_file = os.path.join(args.output_dir, "repeat_annotations_raw.tsv")
        all_annotations_df = pd.concat(all_annotations, ignore_index=True)
        all_annotations_df.to_csv(raw_file, sep='\t', index=False)
        logger.info(f"✓ Raw annotations: {raw_file}")
        logger.info(f"  Total annotations: {len(all_annotations_df):,}")
    
    logger.info("")
    
    # ========================================================================
    # SUMMARY REPORT
    # ========================================================================
    
    logger.info("="*70)
    logger.info("REPEAT PROFILING COMPLETE")
    logger.info("="*70)
    logger.info(f"Samples analyzed: {len(samples_dict)}")
    
    # Calculate totals
    total_classes = len(set(
        cls for comp in all_compositions.values() 
        for cls in comp['class'].keys()
    ))
    total_families = len(set(
        fam for comp in all_compositions.values() 
        for fam in comp['family'].keys()
    ))
    
    logger.info(f"Total unique repeat classes: {total_classes}")
    logger.info(f"Total unique repeat families: {total_families}")
    logger.info("")
    
    logger.info("OUTPUT FILES:")
    for filename in sorted(os.listdir(args.output_dir)):
        filepath = os.path.join(args.output_dir, filename)
        size = os.path.getsize(filepath)
        logger.info(f"  - {filename} ({size:,} bytes)")
    logger.info("")
    
    logger.info("NEXT STEP:")
    logger.info("  python scripts/07_repeat_analysis.py")
    logger.info("")
    logger.info("="*70)

if __name__ == "__main__":
    main()