"""Functions for processing and saving results"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_output.ipynb.

# %% auto 0
__all__ = ['ResultsFormatter', 'ResultsSaver']

# %% ../nbs/02_output.ipynb 4
import os
import pandas as pd
import json
from pathlib import Path
from collections import Counter
from typing import Dict, List, Tuple, Set
from .core import OligoCounter

# %% ../nbs/02_output.ipynb 6
class ResultsFormatter:
    """Formats oligo codon counting results for output."""
    
    @staticmethod
    def to_dataframe(results: OligoCounter, oligos: List[str], offset: int = 1) -> pd.DataFrame:
        """Convert results to a pandas DataFrame.
        
        Args:
            results: Dictionary of oligo index to codon counter
            oligos: List of oligo sequences
            offset: Value to add to oligo index in column names (default: 1)
            
        Returns:
            DataFrame with columns for each oligo and rows for each codon
        """
        df_dict = {}
        for i, oligo in enumerate(oligos):
            col_name = f"{i + offset}_{oligo}"
            col_data = results.get(i, Counter())
            df_dict[col_name] = col_data
        
        # Create DataFrame and fill missing values with 0
        df = pd.DataFrame(df_dict).fillna(0)
        
        # Reorder columns based on oligo index
        df = df.reindex(sorted(df.columns, key=lambda x: int(x.split('_')[0])), axis=1)
        
        return df
    
    @staticmethod
    def summarize_results(df: pd.DataFrame) -> Dict:
        """Generate a summary of the results.
        
        Args:
            df: DataFrame containing oligo codon counts
            
        Returns:
            Dictionary with summary statistics
        """
        summary = {}
        
        # Total reads and matched reads
        for col in df.columns:
            total_reads = df[col].sum()
            matched_reads = total_reads - df[col].get('none', 0)
            
            summary[f"{col}_total_reads"] = total_reads
            summary[f"{col}_matched_reads"] = matched_reads
            if total_reads > 0:
                summary[f"{col}_match_rate"] = matched_reads / total_reads
            else:
                summary[f"{col}_match_rate"] = 0
                
            # Top codons
            if matched_reads > 0:
                # Filter out 'none' and get top 5 codons
                codon_counts = df[col].drop('none', errors='ignore')
                top_codons = codon_counts.nlargest(5)
                summary[f"{col}_top_codons"] = top_codons.to_dict()
            else:
                summary[f"{col}_top_codons"] = {}
        
        return summary

# %% ../nbs/02_output.ipynb 8
class ResultsSaver:
    """Saves processing results to files."""
    
    @staticmethod
    def save_csv(df: pd.DataFrame, output_path: str, 
                 filename: str = 'oligo_counts.csv', create_dir: bool = True) -> str:
        """Save results to a CSV file.
        
        Args:
            df: DataFrame containing oligo codon counts
            output_path: Directory to save the file
            filename: Name of the output file (default: 'oligo_counts.csv')
            create_dir: Whether to create the output directory if it doesn't exist (default: True)
            
        Returns:
            Full path to the saved file
        """
        output_dir = Path(output_path)
        if create_dir:
            output_dir.mkdir(exist_ok=True, parents=True)
            
        output_file = output_dir / filename
        df.to_csv(output_file)
        return str(output_file)
    
    @staticmethod
    def save_summary(results_summary: Dict, output_path: str, 
                     filename: str = 'oligo_summary.json', create_dir: bool = True) -> str:
        """Save a summary of results to a JSON file.
        
        Args:
            results_summary: Dictionary containing summary information
            output_path: Directory to save the file
            filename: Name of the output file (default: 'oligo_summary.json')
            create_dir: Whether to create the output directory if it doesn't exist (default: True)
            
        Returns:
            Full path to the saved file
        """
        output_dir = Path(output_path)
        if create_dir:
            output_dir.mkdir(exist_ok=True, parents=True)
            
        output_file = output_dir / filename
        
        with open(output_file, 'w') as f:
            json.dump(results_summary, f, indent=2)
            
        return str(output_file)
    
    @staticmethod
    def save_text_report(results_summary: Dict, df: pd.DataFrame, output_path: str,
                         filename: str = 'oligo_report.txt', create_dir: bool = True) -> str:
        """Save a human-readable text report of the results.
        
        Args:
            results_summary: Dictionary containing summary information
            df: DataFrame containing oligo codon counts
            output_path: Directory to save the file
            filename: Name of the output file (default: 'oligo_report.txt')
            create_dir: Whether to create the output directory if it doesn't exist (default: True)
            
        Returns:
            Full path to the saved file
        """
        output_dir = Path(output_path)
        if create_dir:
            output_dir.mkdir(exist_ok=True, parents=True)
            
        output_file = output_dir / filename
        
        with open(output_file, 'w') as f:
            f.write("=== OligoSeeker Results Report ===\n\n")
            
            # Write overall stats
            f.write("Overall Statistics:\n")
            f.write(f"Total oligos processed: {len(df.columns)}\n")
            f.write("\n")
            
            # Write per-oligo stats
            f.write("=== Per-Oligo Statistics ===\n\n")
            
            for col in df.columns:
                f.write(f"Oligo: {col}\n")
                f.write(f"  Total reads: {results_summary.get(f'{col}_total_reads', 0)}\n")
                f.write(f"  Matched reads: {results_summary.get(f'{col}_matched_reads', 0)}\n")
                
                match_rate = results_summary.get(f'{col}_match_rate', 0) * 100
                f.write(f"  Match rate: {match_rate:.2f}%\n")
                
                # Top codons
                top_codons = results_summary.get(f'{col}_top_codons', {})
                if top_codons:
                    f.write("  Top codons:\n")
                    for codon, count in top_codons.items():
                        f.write(f"    {codon}: {count}\n")
                
                f.write("\n")
                
            # Write a note about other output files
            f.write("\nNote: More detailed results are available in the CSV, Excel, and JSON output files.\n")
            
        return str(output_file)
