# ------------------------------
# Imports & Dependencies
# ------------------------------
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.messages import HumanMessage, AIMessage
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langgraph.graph import END, StateGraph
from typing_extensions import TypedDict, Annotated
from typing import Sequence, Dict, List, Optional, Any
import chromadb
import numpy as np
import os
import streamlit as st
import requests
import hashlib
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from sklearn.metrics.pairwise import cosine_similarity

# ------------------------------
# State Schema Definition
# ------------------------------
class AgentState(TypedDict):
    messages: Annotated[Sequence[AIMessage | HumanMessage], add_messages]
    context: Dict[str, Any]
    metadata: Dict[str, Any]

# ------------------------------
# Configuration
# ------------------------------
class ResearchConfig:
    DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
    CHROMA_PATH = "chroma_db"
    CHUNK_SIZE = 512
    CHUNK_OVERLAP = 64
    MAX_CONCURRENT_REQUESTS = 5
    EMBEDDING_DIMENSIONS = 1536
    
    DOCUMENT_MAP = {
        "CV-Transformer Model": {
            "title": "Hybrid CV-Transformer Architecture",
            "content": """
            Combines CNN feature extraction with transformer attention mechanisms.
            Key equation: $f(x) = \text{Softmax}(\frac{QK^T}{\sqrt{d_k}})V$
            Achieves 98.2% accuracy on ImageNet-1k with 42ms inference speed
            """
        },
        "Quantum ML": {
            "title": "Quantum Machine Learning",
            "content": """
            Quantum-enhanced optimization techniques for ML models.
            $\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t)$
            100x speedup on optimization tasks with 58% energy reduction
            """
        }
    }
    
    ANALYSIS_TEMPLATE = """Analyze these technical documents:
{context}

Respond in MARKDOWN with:
1. **Key Innovations** (mathematical formulations)
2. **Methodologies** (algorithms & architectures)
3. **Empirical Results** (comparative metrics)
4. **Applications** (industry use cases)
5. **Limitations** (theoretical boundaries)

Include LaTeX equations where applicable."""

# Check for Chroma migration
if os.path.exists(ResearchConfig.CHROMA_PATH):
    st.warning("""
    **ChromDB Migration Required**
    Existing Chroma database detected. Run these commands:
    
    ```bash
    pip install chroma-migrate
    chroma-migrate
    ```
    
    Then restart the application.
    """)
    st.stop()

if not ResearchConfig.DEEPSEEK_API_KEY:
    st.error("""**Configuration Required**  
    1. Get DeepSeek API key: [platform.deepseek.com](https://platform.deepseek.com/)  
    2. Set secret: `DEEPSEEK_API_KEY`  
    3. Rebuild deployment""")
    st.stop()

# ------------------------------
# ChromaDB Document Manager (Updated)
# ------------------------------
class QuantumDocumentManager:
    def __init__(self):
        self.client = chromadb.PersistentClient(path=ResearchConfig.CHROMA_PATH)
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-large",
            dimensions=ResearchConfig.EMBEDDING_DIMENSIONS
        )

    def create_collection(self, document_map: Dict[str, Dict[str, str]], collection_name: str) -> Chroma:
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=ResearchConfig.CHUNK_SIZE,
            chunk_overlap=ResearchConfig.CHUNK_OVERLAP,
            separators=["\n\n", "\n", "|||"]
        )
        
        docs = []
        for key, data in document_map.items():
            chunks = splitter.split_text(data["content"])
            for chunk in chunks:
                docs.append(Document(
                    page_content=chunk,
                    metadata={
                        "title": data["title"],
                        "source": collection_name,
                        "hash": hashlib.sha256(chunk.encode()).hexdigest()[:16]
                    }
                ))
        
        return Chroma.from_documents(
            documents=docs,
            embedding=self.embeddings,
            collection_name=collection_name,
            client=self.client,
            ids=[self._document_id(doc.page_content) for doc in docs]
        )
    
    def _document_id(self, content: str) -> str:
        return f"{hashlib.sha256(content.encode()).hexdigest()[:16]}-{int(time.time())}"

# Initialize document system
qdm = QuantumDocumentManager()
research_docs = qdm.create_collection(ResearchConfig.DOCUMENT_MAP, "research_papers")

# ------------------------------
# Retrieval System
# ------------------------------
class ResearchRetriever:
    def __init__(self):
        self.retriever = research_docs.as_retriever(
            search_type="mmr",
            search_kwargs={
                'k': 4,
                'fetch_k': 20,
                'lambda_mult': 0.85
            }
        )
    
    def retrieve(self, query: str) -> List[Document]:
        try:
            docs = self.retriever.invoke(query)
            if len(docs) < 1:
                raise ValueError("No relevant documents found")
            return docs
        except Exception as e:
            st.error(f"Retrieval Error: {str(e)}")
            return []

# ------------------------------
# Analysis Processor
# ------------------------------
class CognitiveProcessor:
    def __init__(self):
        self.executor = ThreadPoolExecutor(max_workers=ResearchConfig.MAX_CONCURRENT_REQUESTS)
    
    def process_query(self, prompt: str) -> Dict:
        futures = [self.executor.submit(self._api_request, prompt) for _ in range(3)]
        return self._best_result([f.result() for f in as_completed(futures)])

    def _api_request(self, prompt: str) -> Dict:
        headers = {
            "Authorization": f"Bearer {ResearchConfig.DEEPSEEK_API_KEY}",
            "Content-Type": "application/json"
        }
        
        try:
            response = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers=headers,
                json={
                    "model": "deepseek-chat",
                    "messages": [{
                        "role": "user",
                        "content": f"Respond as Senior AI Researcher:\n{prompt}"
                    }],
                    "temperature": 0.7,
                    "max_tokens": 1500,
                    "top_p": 0.9
                },
                timeout=45
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
    
    def _best_result(self, results: List[Dict]) -> Dict:
        valid = [r for r in results if "error" not in r]
        if not valid:
            return {"error": "All API requests failed"}
        
        # Select response with most technical content
        contents = [r.get('choices', [{}])[0].get('message', {}).get('content', '') for r in valid]
        tech_scores = [len(re.findall(r"\$.*?\$", c)) for c in contents]
        return valid[np.argmax(tech_scores)]

# ------------------------------
# Workflow Engine
# ------------------------------
class ResearchWorkflow:
    def __init__(self):
        self.retriever = ResearchRetriever()
        self.processor = CognitiveProcessor()
        self.workflow = StateGraph(AgentState)
        self._build_workflow()

    def _build_workflow(self):
        self.workflow.add_node("ingest", self.ingest)
        self.workflow.add_node("retrieve", self.retrieve)
        self.workflow.add_node("analyze", self.analyze)
        self.workflow.add_node("validate", self.validate)
        self.workflow.add_node("refine", self.refine)

        self.workflow.set_entry_point("ingest")
        self.workflow.add_edge("ingest", "retrieve")
        self.workflow.add_edge("retrieve", "analyze")
        self.workflow.add_conditional_edges(
            "analyze",
            self._quality_check,
            {"valid": "validate", "invalid": "refine"}
        )
        self.workflow.add_edge("validate", END)
        self.workflow.add_edge("refine", "retrieve")

        self.app = self.workflow.compile()

    def ingest(self, state: AgentState) -> Dict:
        try:
            query = state["messages"][-1].content
            return {
                "messages": [AIMessage(content="Query ingested")],
                "context": {"query": query},
                "metadata": {"timestamp": datetime.now().isoformat()}
            }
        except Exception as e:
            return self._error_state(f"Ingestion Error: {str(e)}")

    def retrieve(self, state: AgentState) -> Dict:
        try:
            docs = self.retriever.retrieve(state["context"]["query"])
            return {
                "messages": [AIMessage(content=f"Found {len(docs)} relevant papers")],
                "context": {"docs": docs}
            }
        except Exception as e:
            return self._error_state(f"Retrieval Error: {str(e)}")

    def analyze(self, state: AgentState) -> Dict:
        try:
            context = "\n\n".join([
                f"### {doc.metadata['title']}\n{doc.page_content}" 
                for doc in state["context"]["docs"]
            ])
            prompt = ResearchConfig.ANALYSIS_TEMPLATE.format(context=context)
            response = self.processor.process_query(prompt)
            
            if "error" in response:
                raise RuntimeError(response["error"])
                
            content = response['choices'][0]['message']['content']
            self._validate_analysis(content)
            
            return {"messages": [AIMessage(content=content)]}
        except Exception as e:
            return self._error_state(f"Analysis Error: {str(e)}")

    def validate(self, state: AgentState) -> Dict:
        validation_prompt = f"""Validate this technical analysis:
{state["messages"][-1].content}

Check for:
1. Mathematical accuracy
2. Technical depth
3. Logical consistency

Respond with 'VALID' or 'INVALID'"""
        
        response = self.processor.process_query(validation_prompt)
        valid = "VALID" in response.get('choices', [{}])[0].get('message', {}).get('content', '')
        return {
            "messages": [AIMessage(content=f"{state['messages'][-1].content}\n\nValidation: {'✅ Valid' if valid else '❌ Invalid'}")],
            "context": {"valid": valid}
        }

    def refine(self, state: AgentState) -> Dict:
        refinement_prompt = f"""Improve this analysis:
{state["messages"][-1].content}

Focus on:
1. Mathematical precision
2. Technical terminology
3. Empirical references"""
        
        response = self.processor.process_query(refinement_prompt)
        return {"messages": [AIMessage(content=response['choices'][0]['message']['content'])]}

    def _quality_check(self, state: AgentState) -> str:
        return "valid" if state.get("context", {}).get("valid", False) else "invalid"

    def _validate_analysis(self, content: str):
        required_sections = [
            "Key Innovations",
            "Methodologies", 
            "Empirical Results",
            "Applications",
            "Limitations"
        ]
        missing = [s for s in required_sections if f"## {s}" not in content]
        if missing:
            raise ValueError(f"Missing sections: {', '.join(missing)}")
            
        if not re.search(r"\$.*?\$", content):
            raise ValueError("Analysis lacks mathematical notation")

    def _error_state(self, message: str) -> Dict:
        return {
            "messages": [AIMessage(content=f"❌ {message}")],
            "context": {"error": True},
            "metadata": {"status": "error"}
        }

# ------------------------------
# Streamlit Interface
# ------------------------------
class ResearchInterface:
    def __init__(self):
        self.workflow = ResearchWorkflow()
        self._initialize()

    def _initialize(self):
        st.set_page_config(
            page_title="AI Research Assistant",
            layout="wide",
            initial_sidebar_state="expanded"
        )
        self._inject_styles()
        self._build_sidebar()
        self._build_main()

    def _inject_styles(self):
        st.markdown("""
        <style>
        :root {
            --primary: #2ecc71;
            --secondary: #3498db;
            --background: #0a0a0a;
        }
        .stApp {
            background: var(--background);
            color: white;
        }
        .stTextArea textarea {
            background: #1a1a1a !important;
            border: 2px solid var(--secondary) !important;
        }
        code {
            color: var(--primary);
            background: #002200;
            padding: 2px 4px;
        }
        </style>
        """, unsafe_allow_html=True)

    def _build_sidebar(self):
        with st.sidebar:
            st.title("🔬 Research Corpus")
            for key, data in ResearchConfig.DOCUMENT_MAP.items():
                with st.expander(data["title"]):
                    st.markdown(f"```latex\n{data['content']}\n```")
            st.metric("Vector DB Size", len(research_docs.get()['ids']))

    def _build_main(self):
        st.title("🧠 AI Research Analyst")
        query = st.text_area("Research Query:", height=150,
                           placeholder="Enter technical question...")
        
        if st.button("Analyze", type="primary"):
            self._execute_analysis(query)

    def _execute_analysis(self, query: str):
        try:
            with st.spinner("Analyzing research corpus..."):
                result = self.workflow.app.invoke(
                    {"messages": [HumanMessage(content=query)]}
                )
                
                if result.get("context", {}).get("error"):
                    self._show_error(result["context"]["error"])
                else:
                    self._display_result(result)
        except Exception as e:
            self._show_error(str(e))

    def _display_result(self, result):
        with st.expander("Technical Report", expanded=True):
            st.markdown(result["messages"][-1].content)
            
        with st.expander("Source Excerpts", expanded=False):
            for doc in result["context"].get("docs", []):
                st.markdown(f"**{doc.metadata['title']}**")
                st.code(doc.page_content, language='latex')

    def _show_error(self, message):
        st.error(f"""
        ⚠️ Analysis Failed
        {message}
        
        Mitigation Steps:
        1. Simplify query complexity
        2. Check document connections
        3. Verify technical terms
        """)

if __name__ == "__main__":
    ResearchInterface()