#!/usr/bin/env python

# Standard library imports
import os
import sys
import time
import joblib

# Third-party library imports
from loguru import logger
import numpy as np
import torch


# Python 3.9+
from importlib.resources import files

# BioPython imports
from Bio import  SeqIO

# Local application imports
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))
from src import format_data,  predictor
from src import advanced_argparse

__author__ = "Susanna Grigson"
__maintainer__ = "Susanna Grigson"
__license__ = "MIT"
__version__ = "0.1.3"
__email__ = "susie.grigson@flinders.edu.au"
__status__ = "development"

def parse_arguments():
    parser = advanced_argparse.AdvancedArgumentParser(description="Phynteny Transformer")
    
    # Basic arguments (always visible)
    parser.add_basic_argument("infile", help="Input GenBank file", metavar="INFILE")
    parser.add_basic_argument(
        "-o", "--out", 
        type=str, 
        default="phynteny", 
        help="Output directory",
        required=True
    )
    parser.add_basic_argument(
        "--prefix", 
        type=str, 
        default="phynteny", 
        help="Output prefix"
    )
    parser.add_basic_argument(
        "--esm_model",
        default="facebook/esm2_t33_650M_UR50D",
        choices=[
            "facebook/esm2_t48_15B_UR50D",
            "facebook/esm2_t6_8M_UR50D",
            "facebook/esm2_t12_35M_UR50D",
            "facebook/esm2_t30_150M_UR50D",
            "facebook/esm2_t33_650M_UR50D"
        ],
        help="Specify path to ESM model if not the default"
    )
    parser.add_basic_argument(
        "-m", "--models",
        type=str,
        help="Path to the models to use for predictions",
        default=None  # Changed to None so we can handle it more carefully later
    )
    parser.add_basic_argument(
        "-f", "--force", 
        action="store_true", 
        help="Overwrite output directory"
    )
    parser.add_basic_argument(
        "--batch-size", 
        type=int, 
        default=128, 
        help="Batch size for processing data"
    )
    parser.add_basic_argument(
        "--confidence-threshold", 
        type=float, 
        default=0.8, 
        help="Confidence threshold for high-confidence predictions (0.0-1.0)"
    )
    parser.add_basic_argument(
        "--debug", 
        action="store_true", 
        help="Enable verbose debug logging to console"
    )
    parser.add_basic_argument(
        "--save-attention-weights", 
        action="store_true", 
        help="Save attention weight matrices to files during prediction. "
             "Attention weights will be saved as pickle files in the output directory under 'attention_weights/' folder. "
             "Each genome will have its own file containing attention weights from all models."
    )
    parser.add_basic_argument(
        "--show-model-info", 
        action="store_true", 
        help="Display detailed information about the model architecture including parameter counts and exit"
    )
    parser.add_basic_argument(
        "--version", 
        action="version", 
        version=f"%(prog)s {__version__}"
    )
    
    # Advanced arguments (hidden by default)
    parser.add_advanced_argument(
        "--input-dim", 
        type=int, 
        default=1280, 
        help="Input dimension for the model"
    )
    parser.add_advanced_argument(
        "--num-classes", 
        type=int, 
        default=9, 
        help="Number of classes for the model"
    )
    parser.add_advanced_argument(
        "--num-heads", 
        type=int, 
        default=8, 
        help="Number of attention heads for the model"
    )
    parser.add_advanced_argument(
        "--hidden-dim", 
        type=int, 
        default=512, 
        help="Hidden dimension for the model"
    )
    parser.add_advanced_argument(
        "--lstm-hidden-dim", 
        type=int, 
        default=512, 
        help="LSTM hidden dimension for the model"
    )
    parser.add_advanced_argument(
        "--no-lstm", 
        action="store_true", 
        help="Specify if LSTM should not be used in the model"
    )
    parser.add_advanced_argument(
        "--max-len", 
        type=int, 
        default=1500, 
        help="Maximum length for the model"
    )
    parser.add_advanced_argument(
        "--protein-dropout-rate", 
        type=float, 
        default=0.6, 
        help="Dropout rate for protein features"
    )
    parser.add_advanced_argument(
        "--attention", 
        choices=["absolute", "relative", "circular"], 
        default="circular", 
        help="Type of attention mechanism to use"
    )
    parser.add_advanced_argument(
        "--positional-encoding-type", 
        choices=["fourier", "sinusoidal"], 
        default="sinusoidal", 
        help="Type of positional encoding to use"
    )
    parser.add_advanced_argument(
        "--pre-norm", 
        action="store_true", 
        default=False, 
        help="Use pre-normalization in transformer layers instead of post-normalization"
    )
    parser.add_advanced_argument(
        "--progressive-dropout", 
        action="store_true", 
        default=True, 
        help="Enable progressive dropout for protein features"
    )
    parser.add_advanced_argument(
        "--initial-dropout-rate", 
        type=float, 
        default=1.0, 
        help="Initial dropout rate when using progressive dropout"
    )
    parser.add_advanced_argument(
        "--final-dropout-rate", 
        type=float, 
        default=0.4, 
        help="Final dropout rate when using progressive dropout"
    )
    parser.add_advanced_argument(
        "--output-dim", 
        type=int, 
        default=None, 
        help="Output dimension for the model. Defaults to num_classes if not specified"
    )
    parser.add_advanced_argument(
        "--num-layers", 
        type=int, 
        default=2, 
        help="Number of transformer layers"
    )
    parser.add_advanced_argument(
        "--dropout", 
        type=float, 
        default=0.1, 
        help="Dropout rate for the model"
    )
    
    return parser.parse_args()

def main():
    args = parse_arguments()
    
    start_time = time.time()
    # instantiate output directory 
    format_data.instantiate_output_directory(args.out, args.force)
    
    # Configure logging based on debug flag
    if args.debug:
        # When debug mode is enabled:
        # 1. Always log DEBUG level to the log file
        # 2. Also show DEBUG level logs on the console
        logger.remove()  # Remove default handlers
        logger.add(sys.stderr, level="DEBUG")  # Add console handler with DEBUG level
        logger.add(args.out + "/phynteny.log", level="DEBUG")  # Add file handler
        logger.info("Debug mode enabled - showing verbose logs")
    else:
        # Normal mode: DEBUG to file, INFO to console
        logger.remove()  # Remove default handlers
        logger.add(sys.stderr, level="INFO")  # Add console handler with INFO level
        logger.add(args.out + "/phynteny.log", level="DEBUG")  # Add file handler
    
    # Determine the models directory path
    if args.models is None:
        # Try multiple possible locations for models
        potential_paths = []
        
        # Try site-packages location first (where install_models puts them)
        try:
            # Use importlib.resources instead of pkg_resources
            try:
                # For Python 3.9+
                phynteny_utils_dir = str(files('phynteny_utils').joinpath(''))
            except (ImportError, AttributeError):
                # Fallback for older Python or if module not found
                try:
                    import phynteny_utils
                    phynteny_utils_dir = os.path.dirname(phynteny_utils.__file__)
                except (ImportError, AttributeError):
                    phynteny_utils_dir = None
                    
            if phynteny_utils_dir:
                site_packages_path = os.path.join(phynteny_utils_dir, "models")
                potential_paths.append(site_packages_path)
        except Exception as e:
            logger.debug(f"Error finding phynteny_utils package: {e}")
            
        # Try bin directory location (where script might be looking)
        bin_dir = os.path.dirname(sys.executable)
        bin_path = os.path.join(bin_dir, "phynteny_utils", "models")
        potential_paths.append(bin_path)
        
        # Try path relative to the script
        script_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'phynteny_utils', 'models')
        potential_paths.append(script_path)
        
        # Use the first path that exists and contains model files
        for path in potential_paths:
            if os.path.exists(path) and any(f.endswith('.model') for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))):
                args.models = path
                logger.info(f"Found models at: {path}")
                break
    
    # If no models directory was found or specified, show error
    if args.models is None or not os.path.exists(args.models):
        logger.error("Could not find models directory")
        logger.info("Checked the following locations:")
        for path in potential_paths:
            logger.info(f"  - {path}")
        logger.info("You may need to run 'install_models' to download the models.")
        logger.info("Example: install_models -o /path/to/store/models")
        logger.info("Then run this command with: -m /path/to/store/models")
        sys.exit(1)
    
    # Check for model files directly in the directory
    model_files = [f for f in os.listdir(args.models) if f.endswith('.model')]
    if not model_files:
        logger.error(f"No model files found in: {args.models}")
        logger.info("You need to run 'install_models' to download the models.")
        logger.info("Example: install_models -o /path/to/store/models")
        logger.info("Then run this command with: -m /path/to/store/models")
        sys.exit(1)
    
    logger.info(f"Found {len(model_files)} model files in: {args.models}")

    # Verify integer_category.pkl exists - simplified for Python 3.9+
    integer_file_path = None
    potential_int_paths = []
    
    # Use importlib.resources directly since we're on Python 3.9+
    try:
        # Python 3.9+ approach
        phynteny_utils_dir = str(files('phynteny_utils').joinpath(''))
        integer_file_path = os.path.join(phynteny_utils_dir, "integer_category.pkl")
        potential_int_paths.append(integer_file_path)
    except Exception as e:
        logger.debug(f"Error finding phynteny_utils package via importlib.resources: {e}")
    
    # Add fallback paths just to be safe
    # Check in the same directory as models
    models_parent = os.path.dirname(args.models)
    fallback_path = os.path.join(models_parent, "phynteny_utils", "integer_category.pkl")
    potential_int_paths.append(fallback_path)
    
    # Check relative to script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    script_path = os.path.join(script_dir, "phynteny_utils", "integer_category.pkl")
    potential_int_paths.append(script_path)
    
    # Try all potential paths
    for path in potential_int_paths:
        if os.path.exists(path):
            integer_file_path = path
            logger.info(f"Found integer_category.pkl at: {integer_file_path}")
            break
    
    if integer_file_path is None or not os.path.exists(integer_file_path):
        logger.error("integer_category.pkl not found in any expected location:")
        for path in potential_int_paths:
            logger.info(f"  - {path}")
        sys.exit("Critical error: integer_category.pkl is missing. Please ensure it is correctly installed.")
    
    logger.info(f"Using integer_category.pkl from: {integer_file_path}")

    # read in needed info on PHROGS 
    category_dict, phrog_integer_category = format_data.read_annotations_information()

    # Determine whether to use LSTM based on the --no-lstm flag
    use_lstm = not args.no_lstm

    # preamble 
    logger.info(f'Starting Phynteny transformer v{__version__}') 
    logger.info(f'Command executed: in={args.infile}, out={args.out}, esm_model={args.esm_model}, models={args.models}, force={args.force}')
    logger.info(f'Repository homepage is https://github.com/susiegriggo/Phynteny_transformer/tree/main')
    logger.info(f"Written by {__author__}: {__email__}")

    # read in genbank file
    gb_dict = format_data.read_genbank_file(args.infile, phrog_integer_category)
    logger.info("Genbank file read")
    
    # Log initial data dimensions
    num_genomes = len(gb_dict)
    logger.info(f"Number of genomes loaded from GenBank: {num_genomes}")
    
    # Log number of genes in each genome
    for genome_id, genome_data in gb_dict.items():
        if isinstance(genome_data, dict) and 'sequence' in genome_data:
            sequences = genome_data.get('sequence', [])
            #logger.debug(f"Genome {genome_id} has {len(sequences)} genes")
        else:
            logger.warning(f"Genome {genome_id} does not have expected structure")

    # Extract fasta from the genbank files
    logger.info("Extracting fasta from genbank")
    data_keys = list(gb_dict.keys())
    fasta_out = os.path.join(args.out, f"{args.prefix}.fasta")
    records = [gb_dict.get(k).get("sequence") for k in data_keys]
    with open(fasta_out, "w") as output_handle:
        for r in records:
            SeqIO.write(r, output_handle, "fasta")
    output_handle.close()

    # Extract the embeddings from the outputted fasta files
    logger.info("Computing ESM embeddings")
    logger.info("... if step is being slow consider using GPU!")
    embeddings = format_data.extract_embeddings(
        fasta_out, args.out, model_name=args.esm_model, tokens_per_batch=128,
    )

    # Prepare data
    X, y = format_data.prepare_data(embeddings, gb_dict)
    
    # Check data structure alignment
    logger.info(f"Data preparation complete. Checking dimensions:")
    logger.info(f"Number of genomes: {len(gb_dict)}")
    logger.debug(f"Number of genomes in X: {len(X)}")
    logger.debug(f"Number of genomes in y: {len(y)}")
    
    # Check if any genomes were dropped during data preparation
    if len(X) != len(gb_dict) or len(y) != len(gb_dict):
        missing_in_X = set(gb_dict.keys()) - set(X.keys())
        missing_in_y = set(gb_dict.keys()) - set(y.keys())
        logger.warning(f"Some genomes are missing in prepared data. Missing in X: {missing_in_X}, Missing in y: {missing_in_y}")
    
    # Check gene counts match between gb_dict and X/y for each genome
    for genome_id in X.keys():
        if genome_id in gb_dict:
            genome_data = gb_dict[genome_id]
            if isinstance(genome_data, dict) and 'sequence' in genome_data:
                num_genes_gb = len(genome_data['sequence'])
                num_genes_X = X[genome_id].shape[0]
                num_genes_y = y[genome_id].shape[0]
                if num_genes_gb != num_genes_X or num_genes_gb != num_genes_y:
                    logger.warning(f"Gene count mismatch for genome {genome_id}: gb_dict={num_genes_gb}, X={num_genes_X}, y={num_genes_y}")

    # Get the models 
    logger.info("Creating predictor object")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    p = predictor.Predictor(device=device)
    
    logger.info(f"Reading models from directory: {args.models}")
    logger.info(f"Model parameters: input_dim={args.input_dim}, num_classes={args.num_classes}, num_heads={args.num_heads}, "
                f"hidden_dim={args.hidden_dim}, lstm_hidden_dim={args.lstm_hidden_dim}, dropout={args.dropout}, "
                f"use_lstm={use_lstm}, max_len={args.max_len}, protein_dropout_rate={args.protein_dropout_rate}, "
                f"attention={args.attention}, positional_encoding_type={args.positional_encoding_type}, "
                f"pre_norm={args.pre_norm}, progressive_dropout={args.progressive_dropout}, "
                f"initial_dropout_rate={args.initial_dropout_rate}, final_dropout_rate={args.final_dropout_rate}, "
                f"output_dim={args.output_dim}, num_layers={args.num_layers}")
                
    p.read_models_from_directory(
        args.models, 
        input_dim=args.input_dim, 
        num_classes=args.num_classes, 
        num_heads=args.num_heads, 
        hidden_dim=args.hidden_dim, 
        lstm_hidden_dim=args.lstm_hidden_dim, 
        dropout=args.dropout, 
        use_lstm=use_lstm, 
        max_len=args.max_len, 
        protein_dropout_rate=args.protein_dropout_rate, 
        attention=args.attention, 
        positional_encoding_type=args.positional_encoding_type, 
        pre_norm=args.pre_norm, 
        progressive_dropout=args.progressive_dropout, 
        initial_dropout_rate=args.initial_dropout_rate, 
        final_dropout_rate=args.final_dropout_rate, 
        output_dim=args.output_dim,
        num_layers=args.num_layers
    )

    # If --show-model-info flag is set, display model info and exit
    if args.show_model_info:
        logger.info("Displaying model architecture information...")
        model_info = p.show_model_info()
        
        # Also save model info to a file for reference
        import json
        model_info_file = os.path.join(args.out, "model_info.json")
        with open(model_info_file, 'w') as f:
            # Convert to JSON-serializable format
            json_info = {}
            for model_name, info in model_info.items():
                json_info[model_name] = {
                    'type': info['type'],
                    'total_params': info['total_params'],
                    'trainable_params': info['trainable_params'],
                    'param_breakdown': info['param_breakdown']
                }
            json.dump(json_info, f, indent=2)
        
        logger.info(f"Model information saved to: {model_info_file}")
        logger.info("Model information display complete. Exiting.")
        sys.exit(0)

    # Read the calibration models - support both individual and ensemble calibration
    calibration_models_path = os.path.join(args.models, "calibration_models.pkl")
    ensemble_path = os.path.join(args.models, "all_calibration_models.pkl")

    # Check if ensemble calibration models exist
    if os.path.exists(ensemble_path):
        with open(ensemble_path, 'rb') as f:
            calibration_models = joblib.load(f)
        logger.info(f"Loaded ensemble of {len(calibration_models)} isotonic regression calibration models")
        logger.info("USING ENSEMBLE ISOTONIC REGRESSION for confidence calibration")
    # Fall back to single model calibration
    elif os.path.exists(calibration_models_path):
        with open(calibration_models_path, 'rb') as f:
            calibration_models = joblib.load(f)
        logger.info("Loaded single isotonic regression calibration model")
        logger.info("USING ISOTONIC REGRESSION for confidence calibration")
    else:
        logger.error(f"Calibration models file not found at {calibration_models_path} or {ensemble_path}")
        logger.error("The isotonic regression calibration models are required.")
        logger.error("Please ensure you have the latest models installed.")
        sys.exit(1)

    logger.info("Making predictions using batch inference")
    
    # Use the batch inference method
    if args.save_attention_weights:
        logger.info("Attention weight saving is enabled")
        all_scores, attention_weights = p.predict_inference(
            X, y, 
            batch_size=args.batch_size, 
            return_attention_weights=True, 
            output_dir=args.out
        )
        logger.info(f"Saved attention weights for {len(attention_weights)} genomes")
        
        # Create a summary of attention weights
        p.save_attention_weights_summary(attention_weights, args.out)
    else:
        all_scores = p.predict_inference(X, y, batch_size=args.batch_size)
    
    # Convert all_scores format to match what's expected by confidence calculation
    scores = [all_scores[key] for key in all_scores.keys()]
    
    logger.debug(f"Number of valid scores collected: {len(scores)}")
    logger.debug(f"Number of genomes in original data: {len(gb_dict)}")
    logger.debug(f"Number of genomes in X: {len(X)}")
    
    # Verify that we got scores for all genomes
    if len(scores) != len(X):
        logger.warning(f"Score count mismatch: {len(scores)} scores vs {len(X)} genomes in X")
        # Identify which genomes are missing
        genome_ids = list(X.keys())
        if len(genome_ids) > len(scores):
            logger.warning(f"Missing scores for some genomes")

    logger.info("Calculating confidence scores using isotonic regression calibration")
    
    # Always use isotonic regression for confidence calculation
    predictions, confidence_scores, raw_scores, calibrated_scores = p.compute_confidence_isotonic(scores, calibration_models, phrog_integer_category)

    # Check predictions and confidence scores dimensions
    logger.debug(f"Data dimensions after prediction:")
    logger.debug(f"Number of prediction arrays: {len(predictions)}")
    logger.debug(f"Number of score arrays: {len(scores)}")
    logger.debug(f"Number of confidence score arrays: {len(confidence_scores)}")
    
    # Check if predictions match the expected number of genomes and genes
    if len(predictions) != len(scores):
        logger.error(f"Prediction count mismatch: {len(predictions)} prediction arrays vs {len(scores)} score arrays")
    
    # Debug logging for batching and padding issues
    logger.debug("Checking gene predictions against GenBank data - padding may affect results")
    
    # Check gene count for each genome in predictions
    for i, (genome_id, genome_scores) in enumerate(zip(X.keys(), scores)):
        if genome_id in gb_dict:
            genome_data = gb_dict[genome_id]
            if isinstance(genome_data, dict) and 'sequence' in genome_data:
                num_genes_gb = len(genome_data['sequence'])
                num_genes_X = X[genome_id].shape[0]
                num_predictions = len(predictions[i]) if i < len(predictions) else 0
                
                # Log the detailed structure for this genome
                logger.debug(f"Genome {genome_id} - GenBank genes: {num_genes_gb}, X tensor shape: {X[genome_id].shape}, "
                             f"predictions: {num_predictions}")
                
                if num_genes_gb != num_predictions:
                    logger.debug(f"Batch padding detected for genome {genome_id}: actual genes={num_genes_gb}, padded predictions={num_predictions}")
                    
                    # Check if the predictions array contains padding
                    if num_predictions > num_genes_gb:
                        logger.debug(f"Adjusting predictions by removing batch padding to match actual gene count")
                        # Truncate predictions to match actual gene count if needed
                        if i < len(predictions):
                            predictions[i] = predictions[i][:num_genes_gb]
                            # Carefully handle the scores array based on its structure
                            if isinstance(scores[i], list):
                                # If scores[i] is a list that needs truncation (list of gene scores)
                                if len(scores[i]) > num_genes_gb:
                                    scores[i] = scores[i][:num_genes_gb]
                            else:
                                # For array-type scores, only truncate if first dimension matches genes
                                for score_idx in range(len(scores[i])):
                                    if isinstance(scores[i][score_idx], np.ndarray):
                                        # Only truncate if the shape makes sense to truncate
                                        # This avoids the broadcasting error
                                        if (len(scores[i][score_idx].shape) > 0 and 
                                            scores[i][score_idx].shape[0] > num_genes_gb and 
                                            scores[i][score_idx].shape[0] != args.num_classes):
                                            scores[i][score_idx] = scores[i][score_idx][:num_genes_gb]
                            logger.debug(f"Successfully adjusted predictions for {genome_id}: now {len(predictions[i])} genes (removed padding)")
    
    logger.info("Creating SeqRecord objects for GenBank output")
        
    logger.info(f"Writing predictions to genbank file")
    genbank_file = args.out + "/phynteny.gbk"

    logger.info("Creating GenBank and table outputs")
    
    # Create ordered list of genome IDs for consistent indexing
    genome_ids = list(gb_dict.keys())
    logger.info(f"Writing GenBank and table data for {len(genome_ids)} genomes")
    
    # Create the categories map from calibration_models
    if isinstance(calibration_models, list):
        # If calibration_models is a list (ensemble model), use the first model's keys
        logger.info("Creating category map from the first model in the ensemble")
        categories_map = {i: cat_name for i, cat_name in enumerate(calibration_models[0].keys())}
    else:
        # If calibration_models is a dictionary (single model)
        logger.info("Creating category map from the single model")
        categories_map = {i: cat_name for i, cat_name in enumerate(calibration_models.keys())}

    logger.debug(f"Categories map: {categories_map}")
    
    # Log the mapping to help with debugging
    logger.debug(f"Category mapping created: {categories_map}")
    
    # Write predictions to GenBank file
    p.write_genbank(
        gb_dict, 
        genbank_file, 
        predictions, 
        scores, 
        confidence_scores,
        threshold=args.confidence_threshold,
        categories_map=categories_map
    )
    
    # Log final data structure info before generating the table
    logger.debug("Final data structure check before table generation:")
    logger.debug(f"Number of genomes in gb_dict: {len(gb_dict)}")
    logger.debug(f"Number of prediction arrays: {len(predictions)}")
    logger.debug(f"Number of score arrays: {len(scores)}")
    logger.debug(f"Number of confidence score arrays: {len(confidence_scores)}")
    
    # Generate the table
    table_file = args.out + "/phynteny.tsv"
    format_data.generate_table(
        table_file, 
        gb_dict, 
        category_dict, 
        phrog_integer_category,
        predictions=predictions,
        scores=scores,
        confidence_scores=confidence_scores,
        categories_map=categories_map
    )
    
    logger.info("Predictions completed")
    
    # Provide information about attention weights if they were saved
    if args.save_attention_weights:
        attention_dir = os.path.join(args.out, "attention_weights")
        logger.info(f"Attention weights saved to: {attention_dir}")
        logger.info("Attention weight files:")
        logger.info("  - Individual genome files: {genome_id}_attention_weights.pkl")
        logger.info("  - Summary statistics: attention_weights_summary.txt")
        logger.info("Each pickle file contains a list of attention weight matrices, one per model.")
        logger.info("Attention weight shape: [batch_size, num_heads, seq_length, seq_length]")
    
    end_time = time.time()
    elapsed_time = end_time - start_time
    logger.info(f"Total execution time: {elapsed_time:.2f} seconds")

if __name__ == "__main__":
    main()
