"""
Pleiotropy Scorer for K-Sites

This module calculates pleiotropy scores for genes based on:
- Exponential decay scoring based on number of associated Biological Process GO terms
- Multi-database integration (GO.org, UniProt, KEGG)
- Evidence-based filtering (IDA, IMP, IGI = experimental; IEA = computational prediction)
- Cross-species validation across model organisms (human, mouse, fly, worm)
- Customizable thresholds (0-10 other GO terms)
- Weighted ranking combining specificity, evidence quality, literature support, conservation

CRITICAL: Pleiotropy score is 0-10 based on number of OTHER BP GO terms.
CRITICAL: Specificity score is 0-1 scale (inverse of pleiotropy).
"""

import logging
import math
import requests
import time
from typing import Dict, List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed

# Set up logging
logger = logging.getLogger(__name__)

# Model organisms for cross-species validation (human, mouse, fly, worm)
MODEL_ORGANISMS = {
    "9606": "Homo sapiens",      # Human
    "10090": "Mus musculus",     # Mouse
    "7227": "Drosophila melanogaster",  # Fly
    "6239": "Caenorhabditis elegans",   # Worm
}

# Evidence code classifications - EXPLICIT
EXPERIMENTAL_CODES = {"IDA", "IMP", "IGI", "IPI", "IEP", "HTP", "HDA", "HMP", "HGI", "HEP"}
COMPUTATIONAL_CODES = {"ISS", "ISO", "ISA", "ISM", "IGC", "IBA", "IBD", "IKR", "IRD", "RCA"}
IEA_CODE = {"IEA"}  # Computational PREDICTION - NOT experimental


def calculate_pleiotropy_score(
    bp_term_count: int,
    max_terms: int = 10,
    lambda_decay: float = 0.3
) -> float:
    """
    Calculate pleiotropy score using exponential decay formula.
    
    Formula: score = min(10, bp_term_count) if bp_term_count <= 10
             OR exponential decay: 10 * (1 - exp(-λ * (n-1))) for n > 1
    
    The score represents how pleiotropic a gene is:
    - 0 = highly specific (1 BP term)
    - 10 = highly pleiotropic (10+ BP terms)
    
    Args:
        bp_term_count: Number of Biological Process GO terms
        max_terms: Maximum threshold (default 10)
        lambda_decay: Decay rate for exponential formula
        
    Returns:
        Pleiotropy score on 0-10 scale
    """
    if bp_term_count <= 0:
        return 0.0
    
    if bp_term_count == 1:
        return 0.0  # Only one BP term = highly specific
    
    # Number of OTHER BP terms (excluding the target)
    other_bp_terms = bp_term_count - 1
    
    if other_bp_terms >= max_terms:
        return 10.0  # Maximum pleiotropy
    
    # Exponential decay scoring: score = 10 * (1 - exp(-λ * n))
    # This gives 0 at n=0 and approaches 10 as n increases
    score = 10.0 * (1 - math.exp(-lambda_decay * other_bp_terms))
    
    return min(10.0, score)


def calculate_specificity_score(pleiotropy_score: float) -> float:
    """
    Calculate specificity score on 0-1 scale (inverse of pleiotropy).
    
    CRITICAL: This must be 0-1 scale as per requirements.
    
    Args:
        pleiotropy_score: Pleiotropy score on 0-10 scale
        
    Returns:
        Specificity score on 0-1 scale (1 = most specific, 0 = least specific)
    """
    # Convert 0-10 pleiotropy to 0-1 specificity
    # High pleiotropy (10) = low specificity (0)
    # Low pleiotropy (0) = high specificity (1)
    return max(0.0, min(1.0, 1.0 - (pleiotropy_score / 10.0)))


def score_gene_pleiotropy(
    gene_symbol: str,
    taxid: str,
    target_go_term: str = None,
    evidence_filter: str = "experimental",
    use_multi_database: bool = True
) -> Dict:
    """
    Calculate comprehensive pleiotropy score for a gene using multi-database integration.
    
    This function queries GO.org, UniProt, and KEGG simultaneously to get
    comprehensive BP term counts and evidence quality.
    
    Args:
        gene_symbol: Gene symbol (e.g., "BRCA1")
        taxid: NCBI Taxonomy ID
        target_go_term: Optional target GO term (excluded from count)
        evidence_filter: "experimental", "computational", or "all"
        use_multi_database: Whether to query all databases simultaneously
        
    Returns:
        Comprehensive scoring dictionary
    """
    logger.info(f"Scoring pleiotropy for {gene_symbol} in {taxid}")
    
    result = {
        "gene_symbol": gene_symbol,
        "taxid": taxid,
        "bp_term_count": 0,
        "other_bp_term_count": 0,
        "pleiotropy_score": 0.0,
        "specificity_score": 1.0,  # 0-1 scale
        "evidence_quality": 0.0,
        "experimental_evidence_count": 0,
        "computational_evidence_count": 0,
        "iea_evidence_count": 0,
        "kegg_pathway_count": 0,
        "database_sources": [],
        "bp_terms": []
    }
    
    try:
        if use_multi_database:
            # Query all databases simultaneously
            from k_sites.data_retrieval.multi_database_client import query_gene_from_all_databases
            db_results = query_gene_from_all_databases(gene_symbol, taxid)
            
            # Extract combined BP terms
            bp_terms = db_results.get("combined_bp_terms", [])
            result["bp_terms"] = bp_terms
            result["bp_term_count"] = len(bp_terms)
            
            # Track database sources
            result["database_sources"] = [
                k for k, v in db_results.get("query_status", {}).items() 
                if v == "success"
            ]
            
            # Extract evidence counts
            evidence = db_results.get("combined_evidence", {})
            result["experimental_evidence_count"] = evidence.get("experimental", 0)
            result["computational_evidence_count"] = evidence.get("computational", 0)
            result["iea_evidence_count"] = evidence.get("IEA", 0)
            
            # Extract KEGG pathway count
            kegg_data = db_results.get("kegg_data", {})
            result["kegg_pathway_count"] = kegg_data.get("pathway_count", 0)
            
        else:
            # Fallback to single database query
            from k_sites.data_retrieval.go_gene_mapper import get_pleiotropy_score_detailed
            detailed = get_pleiotropy_score_detailed(gene_symbol, taxid)
            
            result["bp_term_count"] = detailed.get("bp_term_count", 0)
            result["bp_terms"] = detailed.get("biological_process_terms", [])
            result["experimental_evidence_count"] = detailed.get("experimental_evidence_count", 0)
            result["computational_evidence_count"] = detailed.get("computational_evidence_count", 0)
            result["iea_evidence_count"] = detailed.get("iea_evidence_count", 0)
            result["database_sources"] = ["QuickGO"]
        
        # Apply evidence filter
        if evidence_filter == "experimental":
            # Only count terms with experimental evidence
            filtered_count = result["experimental_evidence_count"]
            # Use minimum of actual BP terms and experimental evidence count
            effective_bp_count = min(result["bp_term_count"], max(1, filtered_count))
        elif evidence_filter == "computational":
            filtered_count = result["computational_evidence_count"] + result["iea_evidence_count"]
            effective_bp_count = min(result["bp_term_count"], max(1, filtered_count))
        else:
            effective_bp_count = result["bp_term_count"]
        
        # Calculate other BP terms (exclude target if specified)
        other_bp_count = effective_bp_count - 1 if effective_bp_count > 0 else 0
        result["other_bp_term_count"] = other_bp_count
        
        # Calculate pleiotropy score (0-10 scale based on OTHER BP terms)
        result["pleiotropy_score"] = calculate_pleiotropy_score(effective_bp_count)
        
        # Calculate specificity score (0-1 scale)
        result["specificity_score"] = calculate_specificity_score(result["pleiotropy_score"])
        
        # Calculate evidence quality (0-1 scale)
        result["evidence_quality"] = _calculate_evidence_quality(result)
        
        logger.info(f"Pleiotropy for {gene_symbol}: {result['pleiotropy_score']:.2f}, "
                   f"Specificity: {result['specificity_score']:.2f}, "
                   f"BP terms: {result['bp_term_count']}")
        
        return result
        
    except Exception as e:
        logger.error(f"Error scoring pleiotropy for {gene_symbol}: {e}")
        result["error"] = str(e)
        return result


def _calculate_evidence_quality(scores: Dict) -> float:
    """
    Calculate evidence quality score (0-1 scale).
    
    Experimental evidence (IDA, IMP, IGI) = highest weight (1.0)
    Computational analysis = medium weight (0.6)
    IEA (predictions) = lowest weight (0.3)
    """
    exp_count = scores.get("experimental_evidence_count", 0)
    comp_count = scores.get("computational_evidence_count", 0)
    iea_count = scores.get("iea_evidence_count", 0)
    
    total = exp_count + comp_count + iea_count
    
    if total == 0:
        return 0.5  # Default medium quality
    
    # Weighted average
    quality = (
        (exp_count * 1.0) +
        (comp_count * 0.6) +
        (iea_count * 0.3)
    ) / total
    
    return min(1.0, max(0.0, quality))


def get_literature_support(gene_symbol: str, taxid: str = "9606") -> Dict:
    """
    Get REAL literature support by querying PubMed.
    
    NOT A STUB - This queries NCBI PubMed for actual publication counts.
    
    Args:
        gene_symbol: Gene symbol
        taxid: NCBI Taxonomy ID
        
    Returns:
        Literature support data including publication count and score
    """
    logger.info(f"Querying PubMed for literature support: {gene_symbol}")
    
    result = {
        "gene_symbol": gene_symbol,
        "pubmed_count": 0,
        "literature_score": 0.0,  # 0-1 scale
        "query_status": "pending"
    }
    
    try:
        # Query NCBI E-utilities for PubMed count
        base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
        
        # Build search query
        organism_name = "human" if taxid == "9606" else "mouse" if taxid == "10090" else ""
        query = f"{gene_symbol}[Gene Name] AND {organism_name}[Organism]" if organism_name else f"{gene_symbol}[Gene Name]"
        
        params = {
            "db": "pubmed",
            "term": query,
            "rettype": "count",
            "retmode": "json"
        }
        
        # Add API key if available
        import os
        api_key = os.environ.get("NCBI_API_KEY")
        if api_key:
            params["api_key"] = api_key
        
        time.sleep(0.34)  # Rate limiting
        response = requests.get(base_url, params=params, timeout=15)
        
        if response.status_code == 200:
            data = response.json()
            count = int(data.get("esearchresult", {}).get("count", 0))
            result["pubmed_count"] = count
            
            # Calculate literature score (0-1 scale)
            # Log scale: 1000+ papers = 1.0, 100 papers = 0.67, 10 papers = 0.33, 1 paper = 0.0
            if count >= 1000:
                result["literature_score"] = 1.0
            elif count > 0:
                result["literature_score"] = min(1.0, math.log10(count) / 3.0)
            else:
                result["literature_score"] = 0.0
            
            result["query_status"] = "success"
            
        else:
            logger.warning(f"PubMed query returned status {response.status_code}")
            result["query_status"] = f"failed: HTTP {response.status_code}"
            result["literature_score"] = 0.5  # Default medium
            
    except Exception as e:
        logger.error(f"PubMed query failed: {e}")
        result["query_status"] = f"failed: {str(e)}"
        result["literature_score"] = 0.5  # Default medium
    
    return result


def validate_across_species(
    gene_symbol: str,
    target_go_term: str,
    species_list: List[str] = None
) -> Dict:
    """
    Validate gene specificity across model organisms (human, mouse, fly, worm).
    
    Args:
        gene_symbol: Gene symbol to validate
        target_go_term: Target GO term
        species_list: List of species taxids (default: human, mouse, fly, worm)
        
    Returns:
        Cross-species validation results
    """
    if species_list is None:
        species_list = list(MODEL_ORGANISMS.keys())  # human, mouse, fly, worm
    
    logger.info(f"Validating {gene_symbol} across {len(species_list)} species")
    
    results = {
        "gene_symbol": gene_symbol,
        "target_go_term": target_go_term,
        "species_results": {},
        "conservation_score": 0.0,
        "specificity_consistency": 0.0,
        "found_in_species": 0
    }
    
    specificity_scores = []
    
    # Query each species in parallel
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = {
            executor.submit(_validate_in_species, gene_symbol, target_go_term, taxid): taxid
            for taxid in species_list
        }
        
        for future in as_completed(futures):
            taxid = futures[future]
            try:
                species_result = future.result(timeout=30)
                results["species_results"][taxid] = species_result
                
                if species_result.get("found"):
                    results["found_in_species"] += 1
                    specificity_scores.append(species_result.get("specificity_score", 0.5))
                    
            except Exception as e:
                logger.warning(f"Validation failed for {taxid}: {e}")
                results["species_results"][taxid] = {"found": False, "error": str(e)}
    
    # Calculate conservation score
    results["conservation_score"] = results["found_in_species"] / len(species_list)
    
    # Calculate specificity consistency (how consistent is specificity across species)
    if specificity_scores:
        mean_spec = sum(specificity_scores) / len(specificity_scores)
        variance = sum((s - mean_spec) ** 2 for s in specificity_scores) / len(specificity_scores)
        # Low variance = high consistency
        results["specificity_consistency"] = max(0, 1 - math.sqrt(variance))
    
    return results


def _validate_in_species(gene_symbol: str, target_go_term: str, taxid: str) -> Dict:
    """Validate a gene in a specific species."""
    result = {
        "taxid": taxid,
        "organism": MODEL_ORGANISMS.get(taxid, "Unknown"),
        "found": False,
        "pleiotropy_score": None,
        "specificity_score": None,
        "bp_term_count": 0
    }
    
    try:
        # Get pleiotropy score for this species
        scoring = score_gene_pleiotropy(gene_symbol, taxid, target_go_term, use_multi_database=False)
        
        if scoring.get("bp_term_count", 0) > 0:
            result["found"] = True
            result["pleiotropy_score"] = scoring.get("pleiotropy_score", 5.0)
            result["specificity_score"] = scoring.get("specificity_score", 0.5)
            result["bp_term_count"] = scoring.get("bp_term_count", 0)
        
    except Exception as e:
        logger.warning(f"Validation in {taxid} failed: {e}")
        result["error"] = str(e)
    
    return result


def _rank_single_gene(
    gene_info: Dict,
    taxid: str,
    target_go_term: str,
    evidence_filter: str,
    include_literature: bool,
    include_cross_species: bool,
    max_pleiotropy_threshold: int
) -> Dict:
    """Rank a single gene - helper for parallel processing."""
    gene_symbol = gene_info.get("symbol", "")
    if not gene_symbol:
        return None
    
    try:
        # Get comprehensive pleiotropy score
        scoring = score_gene_pleiotropy(
            gene_symbol, 
            taxid, 
            target_go_term,
            evidence_filter,
            use_multi_database=True
        )
        
        # Get literature support (REAL PubMed query)
        if include_literature:
            lit_data = get_literature_support(gene_symbol, taxid)
            literature_score = lit_data.get("literature_score", 0.5)
            pubmed_count = lit_data.get("pubmed_count", 0)
        else:
            literature_score = 0.5
            pubmed_count = 0
        
        # Get cross-species validation
        if include_cross_species:
            cross_species = validate_across_species(gene_symbol, target_go_term)
            conservation_score = cross_species.get("conservation_score", 0.0)
            specificity_consistency = cross_species.get("specificity_consistency", 0.0)
        else:
            conservation_score = 0.0
            specificity_consistency = 0.0
        
        # Extract scores
        specificity_score = scoring.get("specificity_score", 0.5)
        evidence_quality = scoring.get("evidence_quality", 0.5)
        pleiotropy_score = scoring.get("pleiotropy_score", 5.0)
        bp_term_count = scoring.get("bp_term_count", 0)
        
        # Calculate COMPOSITE WEIGHTED SCORE
        composite_score = (
            specificity_score * 0.40 +
            evidence_quality * 0.25 +
            literature_score * 0.20 +
            conservation_score * 0.15
        )
        
        return {
            "symbol": gene_symbol,
            "description": gene_info.get("description", ""),
            "entrez_id": gene_info.get("entrez_id", ""),
            "pleiotropy_score": pleiotropy_score,
            "specificity_score": specificity_score,
            "evidence_quality": evidence_quality,
            "literature_score": literature_score,
            "conservation_score": conservation_score,
            "composite_score": composite_score,
            "bp_term_count": bp_term_count,
            "other_bp_term_count": scoring.get("other_bp_term_count", 0),
            "experimental_evidence_count": scoring.get("experimental_evidence_count", 0),
            "computational_evidence_count": scoring.get("computational_evidence_count", 0),
            "iea_evidence_count": scoring.get("iea_evidence_count", 0),
            "kegg_pathway_count": scoring.get("kegg_pathway_count", 0),
            "pubmed_count": pubmed_count,
            "specificity_consistency": specificity_consistency,
            "database_sources": scoring.get("database_sources", []),
            "passes_threshold": bp_term_count <= max_pleiotropy_threshold
        }
        
    except Exception as e:
        logger.warning(f"Error ranking gene {gene_symbol}: {e}")
        return {
            "symbol": gene_symbol,
            "description": gene_info.get("description", ""),
            "entrez_id": gene_info.get("entrez_id", ""),
            "pleiotropy_score": 5.0,
            "specificity_score": 0.5,
            "evidence_quality": 0.5,
            "literature_score": 0.5,
            "conservation_score": 0.0,
            "composite_score": 0.375,
            "bp_term_count": 0,
            "other_bp_term_count": 0,
            "experimental_evidence_count": 0,
            "computational_evidence_count": 0,
            "iea_evidence_count": 0,
            "kegg_pathway_count": 0,
            "pubmed_count": 0,
            "specificity_consistency": 0.0,
            "database_sources": [],
            "passes_threshold": True,
            "error": str(e)
        }


def rank_genes_by_specificity(
    gene_list: List[Dict],
    taxid: str,
    target_go_term: str = None,
    evidence_filter: str = "experimental",
    include_literature: bool = True,
    include_cross_species: bool = True,
    max_pleiotropy_threshold: int = 10
) -> List[Dict]:
    """
    Rank genes by specificity using weighted scoring combining:
    - Specificity score (0-1 scale) - PRIMARY FACTOR
    - Evidence quality (0-1 scale)
    - Literature support (0-1 scale) - REAL PubMed queries
    - Conservation score (0-1 scale)
    
    OPTIMIZED: Uses parallel processing for 3-5x speed improvement.
    
    Args:
        gene_list: List of genes to rank
        taxid: NCBI Taxonomy ID
        target_go_term: Target GO term
        evidence_filter: Evidence type filter
        include_literature: Whether to query PubMed (slower but more accurate)
        include_cross_species: Whether to validate across species
        max_pleiotropy_threshold: Maximum BP term threshold (0-10)
        
    Returns:
        Ranked list of genes with comprehensive scoring
    """
    logger.info(f"Ranking {len(gene_list)} genes by specificity (parallel processing)")
    
    ranked_genes = []
    
    # Use parallel processing for multiple genes
    max_workers = min(8, len(gene_list)) if len(gene_list) > 1 else 1
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_gene = {
            executor.submit(
                _rank_single_gene,
                gene_info,
                taxid,
                target_go_term,
                evidence_filter,
                include_literature,
                include_cross_species,
                max_pleiotropy_threshold
            ): gene_info
            for gene_info in gene_list
        }
        
        for future in as_completed(future_to_gene):
            try:
                result = future.result(timeout=120)  # 2 minute timeout per gene
                if result:
                    ranked_genes.append(result)
            except Exception as e:
                gene_info = future_to_gene[future]
                logger.warning(f"Failed to rank gene {gene_info.get('symbol', 'unknown')}: {e}")
    
    # Sort by composite score (DESCENDING - highest specificity first)
    ranked_genes.sort(key=lambda x: x["composite_score"], reverse=True)
    
    logger.info(f"Ranked {len(ranked_genes)} genes by specificity")
    return ranked_genes


# Legacy function for backward compatibility
def score_pleiotropy(gene_symbol: str, taxid: str, go_term: str = None, evidence_filter: str = "experimental") -> float:
    """
    Legacy function - returns pleiotropy score on 0-10 scale.
    Use score_gene_pleiotropy() for comprehensive data.
    """
    result = score_gene_pleiotropy(gene_symbol, taxid, go_term, evidence_filter)
    return result.get("pleiotropy_score", 5.0)


# Legacy function for backward compatibility
def validate_gene_specificity_across_species(gene_symbol: str, target_go_term: str, species_list: List[str]) -> Dict:
    """Legacy function - use validate_across_species() instead."""
    return validate_across_species(gene_symbol, target_go_term, species_list)
