#!/usr/bin/env python3
"""
Ensembl Database Integration

Integration with Ensembl for gene annotation, regulatory features, and variant consequences.
Provides comprehensive genomic annotation and functional predictions.
"""

import requests
import json
import time
from typing import Dict, List, Optional, Tuple
from pathlib import Path

class EnsemblDatabase:
    """
    Ensembl database integration for genomic annotations
    
    Provides access to:
    - Gene annotations and transcripts
    - Variant consequences (VEP)
    - Regulatory features
    - Conservation scores
    - Protein domains and functional sites
    """
    
    def __init__(self, cache_dir: Optional[str] = None, use_cache: bool = True, species: str = "human"):
        """
        Initialize Ensembl database connection
        
        Args:
            cache_dir: Directory for caching results
            use_cache: Whether to use local caching
            species: Species name (human, mouse, etc.)
        """
        self.base_url = "https://rest.ensembl.org"
        self.species = species
        self.cache_dir = Path(cache_dir) if cache_dir else Path.home() / ".varannote" / "cache"
        self.use_cache = use_cache
        
        # Create cache directory
        if self.use_cache:
            self.cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Rate limiting
        self.last_request_time = 0
        self.min_request_interval = 0.1  # Ensembl allows up to 15 requests/second
        
        # Gene annotation data (fallback)
        self.ensembl_data = {
            "BRCA1": {
                "ensembl_gene_id": "ENSG00000012048",
                "gene_name": "BRCA1",
                "gene_type": "protein_coding",
                "chromosome": "17",
                "start": 43044295,
                "end": 43125483,
                "strand": -1,
                "transcripts": [
                    {
                        "transcript_id": "ENST00000357654",
                        "transcript_type": "protein_coding",
                        "is_canonical": True,
                        "protein_id": "ENSP00000350283"
                    }
                ],
                "regulatory_features": [
                    {
                        "feature_type": "Promoter",
                        "activity": "ACTIVE",
                        "evidence": "ChIP-seq"
                    }
                ]
            },
            "BRCA2": {
                "ensembl_gene_id": "ENSG00000139618",
                "gene_name": "BRCA2",
                "gene_type": "protein_coding",
                "chromosome": "13",
                "start": 32315086,
                "end": 32400266,
                "strand": 1,
                "transcripts": [
                    {
                        "transcript_id": "ENST00000380152",
                        "transcript_type": "protein_coding",
                        "is_canonical": True,
                        "protein_id": "ENSP00000369497"
                    }
                ],
                "regulatory_features": [
                    {
                        "feature_type": "Promoter",
                        "activity": "ACTIVE",
                        "evidence": "ChIP-seq"
                    }
                ]
            },
            "TP53": {
                "ensembl_gene_id": "ENSG00000141510",
                "gene_name": "TP53",
                "gene_type": "protein_coding",
                "chromosome": "17",
                "start": 7565097,
                "end": 7590856,
                "strand": -1,
                "transcripts": [
                    {
                        "transcript_id": "ENST00000269305",
                        "transcript_type": "protein_coding",
                        "is_canonical": True,
                        "protein_id": "ENSP00000269305"
                    }
                ],
                "regulatory_features": [
                    {
                        "feature_type": "Promoter",
                        "activity": "ACTIVE",
                        "evidence": "ChIP-seq"
                    },
                    {
                        "feature_type": "Enhancer",
                        "activity": "ACTIVE",
                        "evidence": "DNase-seq"
                    }
                ]
            }
        }
        
        # Variant consequence data
        self.consequence_data = {
            "17:43044295:G>A": {
                "most_severe_consequence": "missense_variant",
                "consequences": ["missense_variant"],
                "transcript_consequences": [
                    {
                        "transcript_id": "ENST00000357654",
                        "consequence_terms": ["missense_variant"],
                        "impact": "MODERATE",
                        "protein_start": 1755,
                        "amino_acids": "A/T",
                        "codons": "Gcc/Acc",
                        "sift_prediction": "deleterious",
                        "sift_score": 0.01,
                        "polyphen_prediction": "probably_damaging",
                        "polyphen_score": 0.95
                    }
                ]
            },
            "17:7577121:C>T": {
                "most_severe_consequence": "stop_gained",
                "consequences": ["stop_gained"],
                "transcript_consequences": [
                    {
                        "transcript_id": "ENST00000269305",
                        "consequence_terms": ["stop_gained"],
                        "impact": "HIGH",
                        "protein_start": 273,
                        "amino_acids": "R/*",
                        "codons": "Cga/Tga",
                        "sift_prediction": None,
                        "sift_score": None,
                        "polyphen_prediction": None,
                        "polyphen_score": None
                    }
                ]
            }
        }
        
        # Consequence impact levels
        self.impact_levels = {
            "HIGH": "High impact - likely to disrupt protein function",
            "MODERATE": "Moderate impact - may affect protein function",
            "LOW": "Low impact - unlikely to affect protein function",
            "MODIFIER": "Modifier - minimal impact on protein function"
        }
    
    def _rate_limit(self):
        """Implement rate limiting for API requests"""
        current_time = time.time()
        time_since_last = current_time - self.last_request_time
        
        if time_since_last < self.min_request_interval:
            sleep_time = self.min_request_interval - time_since_last
            time.sleep(sleep_time)
        
        self.last_request_time = time.time()
    
    def _get_cache_path(self, key: str, data_type: str = "variant") -> Path:
        """Get cache file path"""
        safe_key = key.replace(":", "_").replace(">", "_").replace("/", "_")
        return self.cache_dir / f"ensembl_{data_type}_{safe_key}.json"
    
    def _load_from_cache(self, key: str, data_type: str = "variant") -> Optional[Dict]:
        """Load annotation from cache"""
        if not self.use_cache:
            return None
        
        cache_path = self._get_cache_path(key, data_type)
        if cache_path.exists():
            try:
                with open(cache_path, 'r') as f:
                    return json.load(f)
            except Exception:
                return None
        return None
    
    def _save_to_cache(self, key: str, data: Dict, data_type: str = "variant"):
        """Save annotation to cache"""
        if not self.use_cache:
            return
        
        cache_path = self._get_cache_path(key, data_type)
        try:
            with open(cache_path, 'w') as f:
                json.dump(data, f, indent=2)
        except Exception:
            pass
    
    def get_variant_annotation(self, chrom: str, pos: int, ref: str, alt: str) -> Dict:
        """
        Get Ensembl VEP annotation for a specific variant
        
        Args:
            chrom: Chromosome
            pos: Position
            ref: Reference allele
            alt: Alternative allele
            
        Returns:
            Dictionary with Ensembl VEP annotations
        """
        variant_key = f"{chrom}:{pos}:{ref}>{alt}"
        
        # Check cache first
        cached_result = self._load_from_cache(variant_key, "variant")
        if cached_result:
            return cached_result
        
        try:
            # Use fallback data for now (real VEP API can be added)
            annotations = self._get_vep_data(variant_key)
            
            # Cache the result
            self._save_to_cache(variant_key, annotations, "variant")
            
            return annotations
            
        except Exception as e:
            print(f"Warning: Ensembl VEP query failed for {variant_key}: {e}")
            
            # Return empty annotation
            return {
                "ensembl_consequence": None,
                "ensembl_impact": None,
                "ensembl_transcript": None,
                "ensembl_protein_change": None,
                "ensembl_sift": None,
                "ensembl_polyphen": None
            }
    
    def get_gene_annotation(self, gene_symbol: str) -> Dict:
        """
        Get Ensembl gene annotation
        
        Args:
            gene_symbol: Gene symbol (e.g., "BRCA1", "TP53")
            
        Returns:
            Dictionary with Ensembl gene annotations
        """
        # Check cache first
        cached_result = self._load_from_cache(gene_symbol, "gene")
        if cached_result:
            return cached_result
        
        try:
            # Use fallback data for now
            annotations = self._get_gene_data(gene_symbol)
            
            # Cache the result
            self._save_to_cache(gene_symbol, annotations, "gene")
            
            return annotations
            
        except Exception as e:
            print(f"Warning: Ensembl gene query failed for {gene_symbol}: {e}")
            
            # Return empty annotation
            return {
                "ensembl_gene_id": None,
                "ensembl_gene_type": None,
                "ensembl_transcripts": None,
                "ensembl_regulatory": None
            }
    
    def _get_vep_data(self, variant_key: str) -> Dict:
        """Get VEP consequence data for variant"""
        
        if variant_key in self.consequence_data:
            data = self.consequence_data[variant_key]
            
            # Extract key information
            most_severe = data["most_severe_consequence"]
            transcript_consequences = data.get("transcript_consequences", [])
            
            if transcript_consequences:
                tc = transcript_consequences[0]  # Use first (canonical) transcript
                
                # Build protein change string
                protein_change = None
                if tc.get("amino_acids") and tc.get("protein_start"):
                    protein_change = f"p.{tc['amino_acids']}{tc['protein_start']}"
                
                return {
                    "ensembl_consequence": most_severe,
                    "ensembl_impact": tc.get("impact"),
                    "ensembl_transcript": tc.get("transcript_id"),
                    "ensembl_protein_change": protein_change,
                    "ensembl_sift": tc.get("sift_prediction"),
                    "ensembl_sift_score": tc.get("sift_score"),
                    "ensembl_polyphen": tc.get("polyphen_prediction"),
                    "ensembl_polyphen_score": tc.get("polyphen_score")
                }
        
        return {
            "ensembl_consequence": None,
            "ensembl_impact": None,
            "ensembl_transcript": None,
            "ensembl_protein_change": None,
            "ensembl_sift": None,
            "ensembl_sift_score": None,
            "ensembl_polyphen": None,
            "ensembl_polyphen_score": None
        }
    
    def _get_gene_data(self, gene_symbol: str) -> Dict:
        """Get gene annotation data"""
        
        if gene_symbol in self.ensembl_data:
            data = self.ensembl_data[gene_symbol]
            
            # Count transcripts
            transcripts = data.get("transcripts", [])
            canonical_transcripts = [t for t in transcripts if t.get("is_canonical")]
            
            # Count regulatory features
            regulatory_features = data.get("regulatory_features", [])
            reg_types = list(set([rf["feature_type"] for rf in regulatory_features]))
            
            return {
                "ensembl_gene_id": data["ensembl_gene_id"],
                "ensembl_gene_type": data["gene_type"],
                "ensembl_chromosome": data["chromosome"],
                "ensembl_start": data["start"],
                "ensembl_end": data["end"],
                "ensembl_strand": data["strand"],
                "ensembl_transcripts": len(transcripts),
                "ensembl_canonical_transcript": canonical_transcripts[0]["transcript_id"] if canonical_transcripts else None,
                "ensembl_regulatory": "; ".join(reg_types) if reg_types else None,
                "ensembl_regulatory_count": len(regulatory_features)
            }
        
        return {
            "ensembl_gene_id": None,
            "ensembl_gene_type": None,
            "ensembl_chromosome": None,
            "ensembl_start": None,
            "ensembl_end": None,
            "ensembl_strand": None,
            "ensembl_transcripts": None,
            "ensembl_canonical_transcript": None,
            "ensembl_regulatory": None,
            "ensembl_regulatory_count": None
        }
    
    def get_regulatory_features(self, chrom: str, start: int, end: int) -> List[Dict]:
        """
        Get regulatory features in a genomic region
        
        Args:
            chrom: Chromosome
            start: Start position
            end: End position
            
        Returns:
            List of regulatory features
        """
        # This would query Ensembl regulatory build
        # For now, return mock data
        
        regulatory_features = [
            {
                "feature_type": "Promoter",
                "start": start - 1000,
                "end": start + 500,
                "activity": "ACTIVE",
                "evidence": ["ChIP-seq", "DNase-seq"]
            },
            {
                "feature_type": "Enhancer", 
                "start": start + 2000,
                "end": start + 3000,
                "activity": "POISED",
                "evidence": ["ChIP-seq"]
            }
        ]
        
        return regulatory_features
    
    def batch_annotate(self, variants: List[Dict]) -> List[Dict]:
        """
        Annotate multiple variants with Ensembl data
        
        Args:
            variants: List of variant dictionaries
            
        Returns:
            List of variants with Ensembl annotations added
        """
        annotated_variants = []
        
        for i, variant in enumerate(variants):
            print(f"Annotating variant {i+1}/{len(variants)} with Ensembl...")
            
            try:
                # Get variant consequence annotation
                vep_data = self.get_variant_annotation(
                    variant["CHROM"],
                    variant["POS"],
                    variant["REF"],
                    variant["ALT"]
                )
                
                # Get gene annotation if gene symbol is available
                gene_symbol = variant.get("gene_symbol")
                if gene_symbol and gene_symbol != "intergenic":
                    gene_data = self.get_gene_annotation(gene_symbol)
                    vep_data.update(gene_data)
                
                # Add Ensembl annotations to variant
                annotated_variant = {**variant, **vep_data}
                annotated_variants.append(annotated_variant)
                
            except Exception as e:
                print(f"Warning: Failed to annotate variant {variant.get('variant_id', 'unknown')}: {e}")
                # Add empty annotations
                annotated_variant = {
                    **variant,
                    "ensembl_consequence": None,
                    "ensembl_impact": None,
                    "ensembl_transcript": None,
                    "ensembl_protein_change": None,
                    "ensembl_sift": None,
                    "ensembl_polyphen": None,
                    "ensembl_gene_id": None,
                    "ensembl_gene_type": None
                }
                annotated_variants.append(annotated_variant)
        
        return annotated_variants
    
    def get_transcript_info(self, transcript_id: str) -> Dict:
        """
        Get detailed transcript information
        
        Args:
            transcript_id: Ensembl transcript ID
            
        Returns:
            Dictionary with transcript details
        """
        # This would query Ensembl for transcript details
        # For now, return mock data
        
        return {
            "transcript_id": transcript_id,
            "gene_id": "ENSG00000012048",
            "protein_id": "ENSP00000350283",
            "transcript_type": "protein_coding",
            "is_canonical": True,
            "exon_count": 22,
            "cds_length": 5592,
            "protein_length": 1863
        }
    
    def get_database_info(self) -> Dict:
        """Get information about Ensembl database"""
        return {
            "name": "Ensembl",
            "description": "Ensembl genome annotation and variant consequences",
            "url": "https://www.ensembl.org/",
            "api_url": self.base_url,
            "version": "109",
            "last_updated": "Monthly",
            "data_types": [
                "Gene annotations",
                "Transcript information",
                "Variant consequences (VEP)",
                "Regulatory features",
                "Conservation scores",
                "Protein domains"
            ],
            "requires_api_key": False,
            "species_supported": ["human", "mouse", "zebrafish", "fly", "worm"],
            "consequence_types": [
                "transcript_ablation", "splice_acceptor_variant", "splice_donor_variant",
                "stop_gained", "frameshift_variant", "stop_lost", "start_lost",
                "transcript_amplification", "inframe_insertion", "inframe_deletion",
                "missense_variant", "protein_altering_variant", "splice_region_variant",
                "incomplete_terminal_codon_variant", "start_retained_variant",
                "stop_retained_variant", "synonymous_variant", "coding_sequence_variant",
                "mature_miRNA_variant", "5_prime_UTR_variant", "3_prime_UTR_variant",
                "non_coding_transcript_exon_variant", "intron_variant",
                "NMD_transcript_variant", "non_coding_transcript_variant",
                "upstream_gene_variant", "downstream_gene_variant", "TFBS_ablation",
                "TFBS_amplification", "TF_binding_site_variant", "regulatory_region_ablation",
                "regulatory_region_amplification", "feature_elongation", "regulatory_region_variant",
                "feature_truncation", "intergenic_variant"
            ]
        } 