┌─────────────────────────────────────────────────────────────────────────────┐
│                  MEDICAL LITERATURE SEARCH ARCHITECTURE                     │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │ DATA SOURCES                                                        │   │
│  │   PubMed API ──► Clinical Trials ──► Drug Database ──► Guidelines   │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │ MEDICAL NLP                                                         │   │
│  │   Medical NER ─────────► Entity Normalization ─────────► MeSH Map   │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │ KNOWLEDGE LAYER                                                     │   │
│  │   BioMed Embeddings ──────► Knowledge Graph ──────► Vector Store    │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │ INTELLIGENT RETRIEVAL                                               │   │
│  │   Hybrid Search ───────► Evidence Filtering ───────► Citation Rank  │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │ RESPONSE GENERATION                                                 │   │
│  │              ┌──────────────┬──────────────┬──────────────┐         │   │
│  │              ▼              ▼              ▼              ▼         │   │
│  │         RAG Pipeline    Citations    Evidence Grade   Summary       │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

medical-literature/
├── src/
│   ├── __init__.py
│   ├── config.py                # Configuration
│   ├── ingestion/
│   │   ├── __init__.py
│   │   ├── pubmed_client.py     # PubMed API integration
│   │   ├── clinical_trials.py   # ClinicalTrials.gov client
│   │   └── drug_database.py     # Drug interaction data
│   ├── nlp/
│   │   ├── __init__.py
│   │   ├── medical_ner.py       # Medical entity recognition
│   │   ├── mesh_mapper.py       # MeSH term mapping
│   │   └── abbreviations.py     # Medical abbreviation expansion
│   ├── knowledge/
│   │   ├── __init__.py
│   │   ├── embeddings.py        # BioMedical embeddings
│   │   ├── knowledge_graph.py   # Medical knowledge graph
│   │   └── vector_store.py      # Qdrant integration
│   ├── retrieval/
│   │   ├── __init__.py
│   │   ├── hybrid_search.py     # Hybrid retrieval
│   │   └── evidence_filter.py   # Evidence quality filtering
│   ├── generation/
│   │   ├── __init__.py
│   │   ├── rag_pipeline.py      # RAG for Q&A
│   │   └── evidence_grader.py   # GRADE framework
│   └── api/
│       ├── __init__.py
│       └── main.py              # FastAPI application
├── tests/
├── docker-compose.yml
└── requirements.txt

# src/config.py
from pydantic_settings import BaseSettings
from typing import List

class Settings(BaseSettings):
    # API Keys
    openai_api_key: str
    ncbi_api_key: str = ""  # Optional for higher rate limits

    # Models
    embedding_model: str = "pritamdeka/PubMedBERT-mnli-snli-scinli-scitail-mednli-stsb"
    llm_model: str = "gpt-4o"

    # Vector Store
    qdrant_url: str = "http://localhost:6333"
    qdrant_collection: str = "medical_literature"

    # Neo4j for knowledge graph
    neo4j_uri: str = "bolt://localhost:7687"
    neo4j_user: str = "neo4j"
    neo4j_password: str = "password"

    # PubMed settings
    pubmed_batch_size: int = 100
    pubmed_max_results: int = 1000

    # Evidence levels
    evidence_levels: List[str] = [
        "systematic_review",
        "randomized_controlled_trial",
        "cohort_study",
        "case_control",
        "case_report",
        "expert_opinion"
    ]

    class Config:
        env_file = ".env"

settings = Settings()

# src/ingestion/pubmed_client.py
from typing import List, Dict, Any, Optional, AsyncGenerator
from dataclasses import dataclass
import aiohttp
import asyncio
from xml.etree import ElementTree
import re

from ..config import settings

@dataclass
class PubMedArticle:
    pmid: str
    title: str
    abstract: str
    authors: List[str]
    journal: str
    publication_date: str
    doi: Optional[str]
    mesh_terms: List[str]
    keywords: List[str]
    publication_types: List[str]
    citations_count: int = 0

class PubMedClient:
    """Client for PubMed E-utilities API."""

    BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"

    def __init__(self):
        self.api_key = settings.ncbi_api_key

    async def search(
        self,
        query: str,
        max_results: int = 100,
        date_from: str = None,
        date_to: str = None,
        publication_types: List[str] = None
    ) -> List[str]:
        """Search PubMed and return PMIDs."""
        params = {
            "db": "pubmed",
            "term": self._build_query(query, publication_types),
            "retmax": max_results,
            "retmode": "json",
            "sort": "relevance"
        }

        if self.api_key:
            params["api_key"] = self.api_key

        if date_from:
            params["mindate"] = date_from
        if date_to:
            params["maxdate"] = date_to

        async with aiohttp.ClientSession() as session:
            async with session.get(
                f"{self.BASE_URL}/esearch.fcgi",
                params=params
            ) as response:
                data = await response.json()
                return data.get("esearchresult", {}).get("idlist", [])

    async def fetch_articles(
        self,
        pmids: List[str]
    ) -> AsyncGenerator[PubMedArticle, None]:
        """Fetch full article details for PMIDs."""
        # Process in batches
        for i in range(0, len(pmids), settings.pubmed_batch_size):
            batch = pmids[i:i + settings.pubmed_batch_size]

            params = {
                "db": "pubmed",
                "id": ",".join(batch),
                "retmode": "xml"
            }

            if self.api_key:
                params["api_key"] = self.api_key

            async with aiohttp.ClientSession() as session:
                async with session.get(
                    f"{self.BASE_URL}/efetch.fcgi",
                    params=params
                ) as response:
                    xml_content = await response.text()
                    articles = self._parse_xml(xml_content)

                    for article in articles:
                        yield article

            # Rate limiting
            await asyncio.sleep(0.34)  # ~3 requests per second

    def _build_query(
        self,
        query: str,
        publication_types: List[str] = None
    ) -> str:
        """Build PubMed query with filters."""
        parts = [query]

        if publication_types:
            type_filter = " OR ".join([
                f'"{pt}"[Publication Type]'
                for pt in publication_types
            ])
            parts.append(f"({type_filter})")

        return " AND ".join(parts)

    def _parse_xml(self, xml_content: str) -> List[PubMedArticle]:
        """Parse PubMed XML response."""
        articles = []
        root = ElementTree.fromstring(xml_content)

        for article_elem in root.findall(".//PubmedArticle"):
            try:
                article = self._parse_article(article_elem)
                if article:
                    articles.append(article)
            except Exception:
                continue

        return articles

    def _parse_article(self, elem) -> Optional[PubMedArticle]:
        """Parse single article from XML element."""
        medline = elem.find(".//MedlineCitation")
        if medline is None:
            return None

        pmid = medline.findtext(".//PMID", "")
        article = medline.find(".//Article")

        if article is None:
            return None

        # Title
        title = article.findtext(".//ArticleTitle", "")

        # Abstract
        abstract_parts = []
        for abstract_text in article.findall(".//AbstractText"):
            label = abstract_text.get("Label", "")
            text = abstract_text.text or ""
            if label:
                abstract_parts.append(f"{label}: {text}")
            else:
                abstract_parts.append(text)
        abstract = " ".join(abstract_parts)

        # Authors
        authors = []
        for author in article.findall(".//Author"):
            last_name = author.findtext("LastName", "")
            first_name = author.findtext("ForeName", "")
            if last_name:
                authors.append(f"{last_name}, {first_name}".strip(", "))

        # Journal
        journal = article.findtext(".//Journal/Title", "")

        # Publication date
        pub_date = article.find(".//PubDate")
        if pub_date is not None:
            year = pub_date.findtext("Year", "")
            month = pub_date.findtext("Month", "")
            day = pub_date.findtext("Day", "")
            publication_date = f"{year}-{month}-{day}".strip("-")
        else:
            publication_date = ""

        # DOI
        doi = None
        for article_id in elem.findall(".//ArticleId"):
            if article_id.get("IdType") == "doi":
                doi = article_id.text

        # MeSH terms
        mesh_terms = [
            mesh.findtext("DescriptorName", "")
            for mesh in medline.findall(".//MeshHeading")
        ]

        # Keywords
        keywords = [
            kw.text for kw in medline.findall(".//Keyword")
            if kw.text
        ]

        # Publication types
        pub_types = [
            pt.text for pt in article.findall(".//PublicationType")
            if pt.text
        ]

        return PubMedArticle(
            pmid=pmid,
            title=title,
            abstract=abstract,
            authors=authors[:10],  # Limit authors
            journal=journal,
            publication_date=publication_date,
            doi=doi,
            mesh_terms=mesh_terms,
            keywords=keywords,
            publication_types=pub_types
        )

┌─────────────────────────────────────────────────────────────┐
│ PUBMED SEARCH FLOW                                          │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Query ───► esearch.fcgi ───► Returns PMIDs (list of IDs)   │
│                                     │                       │
│                                     ▼                       │
│             efetch.fcgi ◄───────────┘                       │
│                  │                                          │
│                  ▼                                          │
│         Full Article XML (title, abstract, MeSH, authors)   │
│                                                             │
│  RATE LIMITS:                                               │
│  • Without API key: 3 requests/second                       │
│  • With API key: 10 requests/second                         │
│  • Batch up to 100 PMIDs per efetch call                    │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/nlp/medical_ner.py
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass
import spacy
from scispacy.linking import EntityLinker

@dataclass
class MedicalEntity:
    text: str
    label: str  # DISEASE, DRUG, GENE, etc.
    start: int
    end: int
    cui: str = None  # UMLS Concept Unique Identifier
    canonical_name: str = None
    confidence: float = 1.0

class MedicalNER:
    """Medical Named Entity Recognition using scispaCy."""

    def __init__(self):
        # Load biomedical NER model
        self.nlp = spacy.load("en_core_sci_lg")

        # Add UMLS entity linker
        self.nlp.add_pipe(
            "scispacy_linker",
            config={
                "resolve_abbreviations": True,
                "linker_name": "umls"
            }
        )

    def extract_entities(self, text: str) -> List[MedicalEntity]:
        """Extract medical entities from text."""
        doc = self.nlp(text)
        entities = []

        for ent in doc.ents:
            # Get UMLS linking info
            cui = None
            canonical_name = None
            confidence = 1.0

            if hasattr(ent._, "kb_ents") and ent._.kb_ents:
                top_match = ent._.kb_ents[0]
                cui = top_match[0]
                confidence = top_match[1]

                # Get canonical name from UMLS
                linker = self.nlp.get_pipe("scispacy_linker")
                if cui in linker.kb.cui_to_entity:
                    canonical_name = linker.kb.cui_to_entity[cui].canonical_name

            entities.append(MedicalEntity(
                text=ent.text,
                label=ent.label_,
                start=ent.start_char,
                end=ent.end_char,
                cui=cui,
                canonical_name=canonical_name or ent.text,
                confidence=confidence
            ))

        return entities

    def extract_drug_mentions(self, text: str) -> List[MedicalEntity]:
        """Extract drug/medication mentions."""
        entities = self.extract_entities(text)
        return [e for e in entities if self._is_drug_entity(e)]

    def extract_disease_mentions(self, text: str) -> List[MedicalEntity]:
        """Extract disease/condition mentions."""
        entities = self.extract_entities(text)
        return [e for e in entities if self._is_disease_entity(e)]

    def _is_drug_entity(self, entity: MedicalEntity) -> bool:
        """Check if entity is a drug/medication."""
        drug_labels = {"CHEMICAL", "DRUG"}
        return entity.label in drug_labels

    def _is_disease_entity(self, entity: MedicalEntity) -> bool:
        """Check if entity is a disease/condition."""
        disease_labels = {"DISEASE", "DISORDER"}
        return entity.label in disease_labels


# src/nlp/mesh_mapper.py
from typing import List, Dict, Any, Optional
import aiohttp

class MeSHMapper:
    """Map terms to MeSH (Medical Subject Headings) vocabulary."""

    MESH_API = "https://id.nlm.nih.gov/mesh/lookup/descriptor"

    async def map_to_mesh(self, term: str) -> Optional[Dict[str, Any]]:
        """Map a term to MeSH descriptor."""
        params = {
            "label": term,
            "match": "contains",
            "limit": 5
        }

        async with aiohttp.ClientSession() as session:
            async with session.get(self.MESH_API, params=params) as response:
                if response.status == 200:
                    results = await response.json()
                    if results:
                        return {
                            "mesh_id": results[0].get("resource", "").split("/")[-1],
                            "label": results[0].get("label", term),
                            "tree_numbers": results[0].get("treeNumber", [])
                        }
        return None

    async def expand_query(self, query: str, entities: List[MedicalEntity]) -> str:
        """Expand query with MeSH terms for better recall."""
        expanded_terms = [query]

        for entity in entities[:5]:  # Limit expansions
            mesh_info = await self.map_to_mesh(entity.canonical_name or entity.text)
            if mesh_info:
                # Add MeSH term to query
                expanded_terms.append(f'"{mesh_info["label"]}"[MeSH Terms]')

        return " OR ".join(expanded_terms)

┌─────────────────────────────────────────────────────────────┐
│ ENTITY NORMALIZATION                                        │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Input: "heart attack"                                      │
│  ├─ Extracted: "heart attack" (DISEASE)                     │
│  ├─ UMLS CUI: C0027051                                      │
│  └─ Canonical: "Myocardial Infarction"                      │
│                                                             │
│  Input: "MI"                                                │
│  ├─ Extracted: "MI" (DISEASE)                               │
│  ├─ UMLS CUI: C0027051  ← SAME CUI!                         │
│  └─ Canonical: "Myocardial Infarction"                      │
│                                                             │
│  Result: Both queries find the same papers                  │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/ingestion/drug_database.py
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from enum import Enum
import aiohttp

class InteractionSeverity(Enum):
    CONTRAINDICATED = "contraindicated"
    SEVERE = "severe"
    MODERATE = "moderate"
    MINOR = "minor"
    UNKNOWN = "unknown"

@dataclass
class DrugInteraction:
    drug_a: str
    drug_b: str
    severity: InteractionSeverity
    description: str
    mechanism: str
    clinical_effects: List[str]
    management: str
    references: List[str]

class DrugInteractionChecker:
    """Check for drug-drug interactions using multiple sources."""

    def __init__(self):
        # In production, integrate with:
        # - DrugBank API
        # - RxNorm
        # - FDA Drug Interaction Database
        self.interaction_cache = {}

    async def check_interactions(
        self,
        drugs: List[str]
    ) -> List[DrugInteraction]:
        """Check interactions between multiple drugs."""
        interactions = []

        # Check all pairs
        for i, drug_a in enumerate(drugs):
            for drug_b in drugs[i+1:]:
                interaction = await self._check_pair(drug_a, drug_b)
                if interaction:
                    interactions.append(interaction)

        # Sort by severity
        severity_order = {
            InteractionSeverity.CONTRAINDICATED: 0,
            InteractionSeverity.SEVERE: 1,
            InteractionSeverity.MODERATE: 2,
            InteractionSeverity.MINOR: 3,
            InteractionSeverity.UNKNOWN: 4
        }
        interactions.sort(key=lambda x: severity_order[x.severity])

        return interactions

    async def _check_pair(
        self,
        drug_a: str,
        drug_b: str
    ) -> Optional[DrugInteraction]:
        """Check interaction between two drugs."""
        cache_key = tuple(sorted([drug_a.lower(), drug_b.lower()]))

        if cache_key in self.interaction_cache:
            return self.interaction_cache[cache_key]

        # Query interaction database (mock implementation)
        interaction = await self._query_interaction_db(drug_a, drug_b)
        self.interaction_cache[cache_key] = interaction

        return interaction

    async def _query_interaction_db(
        self,
        drug_a: str,
        drug_b: str
    ) -> Optional[DrugInteraction]:
        """Query drug interaction database."""
        # In production, integrate with actual drug databases
        # This is a simplified example

        # Known severe interactions (example data)
        known_interactions = {
            ("warfarin", "aspirin"): DrugInteraction(
                drug_a="Warfarin",
                drug_b="Aspirin",
                severity=InteractionSeverity.SEVERE,
                description="Increased risk of bleeding when combined",
                mechanism="Both drugs affect hemostasis through different mechanisms",
                clinical_effects=["Increased bleeding risk", "Prolonged INR"],
                management="Avoid combination if possible. Monitor INR closely if necessary.",
                references=["PMID:12345678"]
            ),
            ("metformin", "contrast"): DrugInteraction(
                drug_a="Metformin",
                drug_b="Iodinated Contrast",
                severity=InteractionSeverity.SEVERE,
                description="Risk of lactic acidosis with contrast media",
                mechanism="Contrast may cause acute kidney injury, reducing metformin clearance",
                clinical_effects=["Lactic acidosis", "Acute kidney injury"],
                management="Hold metformin before and 48h after contrast administration",
                references=["PMID:23456789"]
            )
        }

        key = tuple(sorted([drug_a.lower(), drug_b.lower()]))
        return known_interactions.get(key)

# src/generation/evidence_grader.py
from typing import List, Dict, Any
from dataclasses import dataclass
from enum import Enum

class EvidenceLevel(Enum):
    HIGH = "high"          # Systematic reviews, high-quality RCTs
    MODERATE = "moderate"  # Lower-quality RCTs, well-designed cohort studies
    LOW = "low"            # Case-control studies, case series
    VERY_LOW = "very_low"  # Expert opinion, case reports

class RecommendationStrength(Enum):
    STRONG = "strong"
    WEAK = "weak"

@dataclass
class EvidenceAssessment:
    level: EvidenceLevel
    recommendation_strength: RecommendationStrength
    study_design: str
    sample_size: int
    quality_factors: Dict[str, bool]
    limitations: List[str]
    summary: str

class EvidenceGrader:
    """Grade evidence quality using GRADE framework."""

    # Publication type to base evidence level mapping
    STUDY_TYPE_LEVELS = {
        "Systematic Review": EvidenceLevel.HIGH,
        "Meta-Analysis": EvidenceLevel.HIGH,
        "Randomized Controlled Trial": EvidenceLevel.HIGH,
        "Controlled Clinical Trial": EvidenceLevel.MODERATE,
        "Cohort Studies": EvidenceLevel.MODERATE,
        "Case-Control Studies": EvidenceLevel.LOW,
        "Case Reports": EvidenceLevel.VERY_LOW,
        "Review": EvidenceLevel.LOW,
        "Editorial": EvidenceLevel.VERY_LOW,
        "Comment": EvidenceLevel.VERY_LOW
    }

    def grade_article(
        self,
        publication_types: List[str],
        abstract: str,
        mesh_terms: List[str]
    ) -> EvidenceAssessment:
        """Grade the evidence level of an article."""
        # Determine study design
        study_design = self._determine_study_design(publication_types)
        base_level = self.STUDY_TYPE_LEVELS.get(study_design, EvidenceLevel.LOW)

        # Extract sample size from abstract (simplified)
        sample_size = self._extract_sample_size(abstract)

        # Assess quality factors
        quality_factors = self._assess_quality(abstract, mesh_terms)

        # Adjust level based on quality
        final_level = self._adjust_level(base_level, quality_factors, sample_size)

        # Identify limitations
        limitations = self._identify_limitations(abstract, quality_factors)

        # Determine recommendation strength
        strength = self._determine_strength(final_level, quality_factors)

        return EvidenceAssessment(
            level=final_level,
            recommendation_strength=strength,
            study_design=study_design,
            sample_size=sample_size,
            quality_factors=quality_factors,
            limitations=limitations,
            summary=self._generate_summary(final_level, study_design, limitations)
        )

    def _determine_study_design(self, publication_types: List[str]) -> str:
        """Determine primary study design."""
        for pub_type in publication_types:
            if pub_type in self.STUDY_TYPE_LEVELS:
                return pub_type
        return "Unknown"

    def _extract_sample_size(self, abstract: str) -> int:
        """Extract sample size from abstract."""
        import re

        # Common patterns for sample size
        patterns = [
            r'n\s*=\s*(\d+)',
            r'(\d+)\s*patients',
            r'(\d+)\s*participants',
            r'(\d+)\s*subjects',
            r'sample size of\s*(\d+)'
        ]

        for pattern in patterns:
            match = re.search(pattern, abstract.lower())
            if match:
                return int(match.group(1))

        return 0

    def _assess_quality(
        self,
        abstract: str,
        mesh_terms: List[str]
    ) -> Dict[str, bool]:
        """Assess quality factors."""
        abstract_lower = abstract.lower()

        return {
            "randomization": "random" in abstract_lower,
            "blinding": any(term in abstract_lower for term in ["blind", "masked"]),
            "placebo_controlled": "placebo" in abstract_lower,
            "intention_to_treat": "intention to treat" in abstract_lower,
            "multicenter": "multicenter" in abstract_lower or "multi-center" in abstract_lower,
            "adequate_followup": any(term in abstract_lower for term in ["follow-up", "followed for"]),
            "low_dropout": "dropout" not in abstract_lower or "low dropout" in abstract_lower
        }

    def _adjust_level(
        self,
        base_level: EvidenceLevel,
        quality_factors: Dict[str, bool],
        sample_size: int
    ) -> EvidenceLevel:
        """Adjust evidence level based on quality."""
        levels = list(EvidenceLevel)
        current_index = levels.index(base_level)

        # Upgrade factors
        upgrades = sum([
            quality_factors.get("multicenter", False),
            sample_size > 500,
            all([quality_factors.get("randomization", False),
                 quality_factors.get("blinding", False)])
        ])

        # Downgrade factors
        downgrades = sum([
            not quality_factors.get("adequate_followup", True),
            sample_size < 50 and sample_size > 0,
            not quality_factors.get("low_dropout", True)
        ])

        # Apply adjustments
        new_index = max(0, min(len(levels) - 1, current_index - upgrades + downgrades))
        return levels[new_index]

    def _identify_limitations(
        self,
        abstract: str,
        quality_factors: Dict[str, bool]
    ) -> List[str]:
        """Identify study limitations."""
        limitations = []
        abstract_lower = abstract.lower()

        if not quality_factors.get("randomization", False):
            limitations.append("Non-randomized design")

        if not quality_factors.get("blinding", False):
            limitations.append("Lack of blinding")

        if "limitation" in abstract_lower:
            limitations.append("Authors report limitations")

        if "small sample" in abstract_lower:
            limitations.append("Small sample size")

        if "single center" in abstract_lower:
            limitations.append("Single center study")

        return limitations

    def _determine_strength(
        self,
        level: EvidenceLevel,
        quality_factors: Dict[str, bool]
    ) -> RecommendationStrength:
        """Determine recommendation strength."""
        high_quality = sum(quality_factors.values()) >= 4

        if level in [EvidenceLevel.HIGH, EvidenceLevel.MODERATE] and high_quality:
            return RecommendationStrength.STRONG
        return RecommendationStrength.WEAK

    def _generate_summary(
        self,
        level: EvidenceLevel,
        study_design: str,
        limitations: List[str]
    ) -> str:
        """Generate evidence summary."""
        summary = f"{level.value.title()} quality evidence from {study_design.lower()}."

        if limitations:
            summary += f" Limitations: {', '.join(limitations[:3])}."

        return summary

┌─────────────────────────────────────────────────────────────┐
│ EVIDENCE PYRAMID                                            │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│                    /\                                       │
│                   /  \  Systematic Reviews                  │
│                  /────\  & Meta-Analyses                    │
│                 /      \                                    │
│                /────────\  Randomized Controlled            │
│               /          \  Trials (RCTs)                   │
│              /────────────\                                 │
│             /              \  Cohort Studies                │
│            /────────────────\                               │
│           /                  \  Case-Control Studies        │
│          /────────────────────\                             │
│         /                      \  Case Reports              │
│        /────────────────────────\                           │
│       /                          \  Expert Opinion          │
│      /____________________________\                         │
│                                                             │
│  HIGH ◄─────────────────────────────────────────► VERY LOW  │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/generation/rag_pipeline.py
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from openai import OpenAI
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer

from ..config import settings
from ..ingestion.pubmed_client import PubMedArticle
from .evidence_grader import EvidenceGrader, EvidenceAssessment

@dataclass
class MedicalAnswer:
    answer: str
    confidence: float
    evidence_level: str
    sources: List[Dict[str, Any]]
    drug_interactions: List[Dict[str, Any]]
    disclaimer: str

class MedicalRAG:
    """RAG pipeline for medical literature Q&A."""

    DISCLAIMER = """This information is for educational purposes only and should not
replace professional medical advice. Always consult with a healthcare provider
for medical decisions."""

    def __init__(self):
        self.openai = OpenAI(api_key=settings.openai_api_key)
        self.qdrant = QdrantClient(url=settings.qdrant_url)

        # BioMedical embeddings
        self.embedder = SentenceTransformer(settings.embedding_model)
        self.evidence_grader = EvidenceGrader()

        self._ensure_collection()

    def _ensure_collection(self):
        """Ensure Qdrant collection exists."""
        from qdrant_client.http.models import Distance, VectorParams

        collections = self.qdrant.get_collections().collections
        exists = any(c.name == settings.qdrant_collection for c in collections)

        if not exists:
            self.qdrant.create_collection(
                collection_name=settings.qdrant_collection,
                vectors_config=VectorParams(
                    size=768,  # PubMedBERT dimension
                    distance=Distance.COSINE
                )
            )

    def index_articles(self, articles: List[PubMedArticle]):
        """Index articles in vector store."""
        from qdrant_client.http.models import PointStruct

        points = []
        for article in articles:
            # Create searchable text
            text = f"{article.title} {article.abstract}"

            # Generate embedding
            embedding = self.embedder.encode(text).tolist()

            # Grade evidence
            evidence = self.evidence_grader.grade_article(
                article.publication_types,
                article.abstract,
                article.mesh_terms
            )

            points.append(PointStruct(
                id=int(article.pmid),
                vector=embedding,
                payload={
                    "pmid": article.pmid,
                    "title": article.title,
                    "abstract": article.abstract,
                    "authors": article.authors,
                    "journal": article.journal,
                    "publication_date": article.publication_date,
                    "doi": article.doi,
                    "mesh_terms": article.mesh_terms,
                    "evidence_level": evidence.level.value,
                    "study_design": evidence.study_design
                }
            ))

        self.qdrant.upsert(
            collection_name=settings.qdrant_collection,
            points=points
        )

    def query(
        self,
        question: str,
        top_k: int = 10,
        min_evidence_level: str = None
    ) -> MedicalAnswer:
        """Answer a medical question with evidence."""
        # Generate query embedding
        query_embedding = self.embedder.encode(question).tolist()

        # Search vector store
        results = self.qdrant.search(
            collection_name=settings.qdrant_collection,
            query_vector=query_embedding,
            limit=top_k
        )

        # Filter by evidence level if specified
        if min_evidence_level:
            level_order = ["high", "moderate", "low", "very_low"]
            min_index = level_order.index(min_evidence_level)
            results = [
                r for r in results
                if level_order.index(r.payload.get("evidence_level", "very_low")) <= min_index
            ]

        if not results:
            return MedicalAnswer(
                answer="No relevant medical literature found for this query.",
                confidence=0.0,
                evidence_level="none",
                sources=[],
                drug_interactions=[],
                disclaimer=self.DISCLAIMER
            )

        # Prepare context
        context = self._prepare_context(results)

        # Generate answer
        answer = self._generate_answer(question, context)

        # Determine overall evidence level
        evidence_levels = [r.payload.get("evidence_level", "very_low") for r in results[:5]]
        overall_evidence = self._aggregate_evidence_level(evidence_levels)

        return MedicalAnswer(
            answer=answer["response"],
            confidence=answer["confidence"],
            evidence_level=overall_evidence,
            sources=[
                {
                    "pmid": r.payload["pmid"],
                    "title": r.payload["title"],
                    "authors": r.payload["authors"][:3],
                    "journal": r.payload["journal"],
                    "evidence_level": r.payload.get("evidence_level"),
                    "relevance_score": r.score
                }
                for r in results[:5]
            ],
            drug_interactions=[],  # Populated separately if drugs detected
            disclaimer=self.DISCLAIMER
        )

    def _prepare_context(self, results) -> str:
        """Prepare context from search results."""
        context_parts = []

        for i, result in enumerate(results[:5]):
            payload = result.payload
            context_parts.append(f"""
[Source {i+1}] PMID: {payload['pmid']}
Title: {payload['title']}
Evidence Level: {payload.get('evidence_level', 'unknown')}
Study Design: {payload.get('study_design', 'unknown')}
Abstract: {payload['abstract'][:1500]}
""")

        return "\n".join(context_parts)

    def _generate_answer(
        self,
        question: str,
        context: str
    ) -> Dict[str, Any]:
        """Generate answer using LLM."""
        prompt = f"""You are a medical research assistant. Answer the question based ONLY on the provided research abstracts.

Research Abstracts:
{context}

Question: {question}

Instructions:
1. Base your answer ONLY on the provided research evidence
2. Cite specific studies using [Source N] format
3. Note the evidence level (high, moderate, low) for key findings
4. Highlight any conflicting evidence
5. If the evidence is insufficient, say so
6. Do NOT provide medical advice or treatment recommendations

Provide your response in JSON format:
{{
    "response": "Your evidence-based answer with citations",
    "confidence": 0.0-1.0,
    "key_findings": ["finding 1", "finding 2"],
    "evidence_gaps": ["gap 1"]
}}"""

        response = self.openai.chat.completions.create(
            model=settings.llm_model,
            messages=[
                {"role": "system", "content": "You are a medical research assistant that provides evidence-based answers."},
                {"role": "user", "content": prompt}
            ],
            response_format={"type": "json_object"},
            temperature=0.1
        )

        import json
        return json.loads(response.choices[0].message.content)

    def _aggregate_evidence_level(self, levels: List[str]) -> str:
        """Aggregate evidence levels from multiple sources."""
        level_scores = {"high": 4, "moderate": 3, "low": 2, "very_low": 1}
        scores = [level_scores.get(level, 1) for level in levels]
        avg_score = sum(scores) / len(scores) if scores else 1

        if avg_score >= 3.5:
            return "high"
        elif avg_score >= 2.5:
            return "moderate"
        elif avg_score >= 1.5:
            return "low"
        return "very_low"

┌─────────────────────────────────────────────────────────────┐
│ EVIDENCE-FILTERED RETRIEVAL                                 │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Query: "Is aspirin effective for heart attack prevention?" │
│                                                             │
│  WITHOUT FILTER:                                            │
│  ├─ Case report from 1985 (n=1)                             │
│  ├─ Editorial with opinion                                  │
│  └─ Small observational study                               │
│                                                             │
│  WITH min_evidence_level="moderate":                        │
│  ├─ Cochrane systematic review (2023)                       │
│  ├─ ASCEND RCT (n=15,480)                                   │
│  └─ ARRIVE trial (n=12,546)                                 │
│                                                             │
│  Result: Reliable, actionable evidence                      │
│                                                             │
└─────────────────────────────────────────────────────────────┘

# src/api/main.py
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any, Optional

from ..config import settings
from ..ingestion.pubmed_client import PubMedClient
from ..ingestion.drug_database import DrugInteractionChecker
from ..nlp.medical_ner import MedicalNER
from ..generation.rag_pipeline import MedicalRAG

app = FastAPI(
    title="Medical Literature Search System",
    description="Evidence-based medical research assistant",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

# Initialize components
pubmed = PubMedClient()
ner = MedicalNER()
drug_checker = DrugInteractionChecker()
rag = MedicalRAG()


class SearchRequest(BaseModel):
    query: str
    max_results: int = 50
    date_from: Optional[str] = None
    study_types: Optional[List[str]] = None

class QuestionRequest(BaseModel):
    question: str
    min_evidence_level: Optional[str] = None

class DrugCheckRequest(BaseModel):
    drugs: List[str]


@app.post("/api/search")
async def search_literature(request: SearchRequest):
    """Search PubMed for relevant articles."""
    # Extract medical entities for query expansion
    entities = ner.extract_entities(request.query)

    # Search PubMed
    pmids = await pubmed.search(
        query=request.query,
        max_results=request.max_results,
        date_from=request.date_from,
        publication_types=request.study_types
    )

    # Fetch articles
    articles = []
    async for article in pubmed.fetch_articles(pmids):
        articles.append({
            "pmid": article.pmid,
            "title": article.title,
            "abstract": article.abstract[:500] + "..." if len(article.abstract) > 500 else article.abstract,
            "authors": article.authors[:5],
            "journal": article.journal,
            "publication_date": article.publication_date,
            "mesh_terms": article.mesh_terms[:10]
        })

        # Index for RAG
        rag.index_articles([article])

    return {
        "query": request.query,
        "entities_detected": [
            {"text": e.text, "type": e.label, "canonical": e.canonical_name}
            for e in entities[:10]
        ],
        "total_results": len(articles),
        "articles": articles
    }


@app.post("/api/question")
async def answer_question(request: QuestionRequest):
    """Answer a medical question using indexed literature."""
    # Extract entities from question
    entities = ner.extract_entities(request.question)

    # Get answer from RAG
    answer = rag.query(
        question=request.question,
        min_evidence_level=request.min_evidence_level
    )

    # Check for drug interactions if drugs mentioned
    drugs = [e.canonical_name for e in entities if ner._is_drug_entity(e)]
    interactions = []
    if len(drugs) >= 2:
        interactions = await drug_checker.check_interactions(drugs)

    return {
        "question": request.question,
        "answer": answer.answer,
        "confidence": answer.confidence,
        "evidence_level": answer.evidence_level,
        "sources": answer.sources,
        "drug_interactions": [
            {
                "drugs": [i.drug_a, i.drug_b],
                "severity": i.severity.value,
                "description": i.description,
                "management": i.management
            }
            for i in interactions
        ],
        "disclaimer": answer.disclaimer
    }


@app.post("/api/drugs/interactions")
async def check_drug_interactions(request: DrugCheckRequest):
    """Check for drug-drug interactions."""
    interactions = await drug_checker.check_interactions(request.drugs)

    return {
        "drugs_checked": request.drugs,
        "interactions_found": len(interactions),
        "interactions": [
            {
                "drug_a": i.drug_a,
                "drug_b": i.drug_b,
                "severity": i.severity.value,
                "description": i.description,
                "mechanism": i.mechanism,
                "clinical_effects": i.clinical_effects,
                "management": i.management
            }
            for i in interactions
        ]
    }


@app.get("/api/health")
async def health_check():
    return {"status": "healthy"}

# docker-compose.yml
version: '3.8'

services:
  api:
    build: .
    ports:
      - "8000:8000"
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - NCBI_API_KEY=${NCBI_API_KEY}
      - QDRANT_URL=http://qdrant:6333
    depends_on:
      - qdrant

  qdrant:
    image: qdrant/qdrant:latest
    ports:
      - "6333:6333"
    volumes:
      - qdrant_data:/qdrant/storage

volumes:
  qdrant_data:

# requirements.txt
fastapi==0.109.0
uvicorn==0.27.0
pydantic==2.5.3
pydantic-settings==2.1.0
openai==1.10.0
qdrant-client==1.7.0
sentence-transformers==2.2.2
aiohttp==3.9.1
scispacy==0.5.3
spacy==3.7.2

import requests

# Search for literature
response = requests.post(
    "http://localhost:8000/api/search",
    json={
        "query": "metformin diabetes type 2 cardiovascular outcomes",
        "max_results": 50,
        "study_types": ["Randomized Controlled Trial", "Meta-Analysis"]
    }
)
results = response.json()
print(f"Found {results['total_results']} articles")

# Ask a question
response = requests.post(
    "http://localhost:8000/api/question",
    json={
        "question": "What is the evidence for metformin reducing cardiovascular mortality in type 2 diabetes?",
        "min_evidence_level": "moderate"
    }
)
answer = response.json()
print(f"Answer: {answer['answer']}")
print(f"Evidence Level: {answer['evidence_level']}")