#!/usr/bin/env python3
"""
GEO Dataset Search Script
Search for suitable gene expression datasets on GEO matching specific criteria.

Search criteria:
- Cancer cell line (MCF7, A549, HeLa, etc.)
- Small molecule treatment (Doxorubicin, Paclitaxel, Bortezomib, etc.)
- Clear treatment vs control comparison
- At least 2 replicates per condition
"""

import requests
import xml.etree.ElementTree as ET
import json
import time
import re
from pathlib import Path

# Search parameters
CELL_LINES = ["MCF7", "MCF-7", "A549", "HeLa", "HCT116", "HepG2"]
COMPOUNDS = ["Doxorubicin", "Paclitaxel", "Bortezomib", "Cisplatin", "Etoposide", "Tamoxifen"]
MIN_REPLICATES = 2

def search_geo_datasets(search_terms, retmax=50):
    """
    Search GEO using NCBI E-utilities API.

    Args:
        search_terms: List of search terms
        retmax: Maximum number of results to return

    Returns:
        List of GSE IDs
    """
    print(f"Searching GEO with terms: {search_terms}")

    # Construct search query
    query = " AND ".join([f"({term})" for term in search_terms])
    query += " AND expression profiling by array[DataSet Type]"

    # E-search to get GSE IDs
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
    params = {
        "db": "gds",
        "term": query,
        "retmax": retmax,
        "retmode": "json"
    }

    print(f"Query: {query}")
    response = requests.get(base_url, params=params)

    if response.status_code != 200:
        print(f"Error: HTTP {response.status_code}")
        return []

    data = response.json()
    id_list = data.get("esearchresult", {}).get("idlist", [])
    print(f"Found {len(id_list)} datasets")

    return id_list

def get_dataset_summary(gds_ids):
    """
    Get summary information for GEO datasets.

    Args:
        gds_ids: List of GEO dataset IDs

    Returns:
        List of dataset summaries
    """
    if not gds_ids:
        return []

    print(f"Fetching summaries for {len(gds_ids)} datasets...")

    # E-summary to get detailed info
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
    params = {
        "db": "gds",
        "id": ",".join(gds_ids[:20]),  # Limit to first 20 to avoid timeout
        "retmode": "json"
    }

    response = requests.get(base_url, params=params)

    if response.status_code != 200:
        print(f"Error: HTTP {response.status_code}")
        return []

    data = response.json()
    results = data.get("result", {})

    summaries = []
    for gds_id in gds_ids[:20]:
        if gds_id in results:
            summary = results[gds_id]
            summaries.append(summary)

    return summaries

def parse_gse_from_summary(summary):
    """Extract GSE ID from dataset summary."""
    # Check in accession field
    accession = summary.get("accession", "")
    if accession.startswith("GSE"):
        return accession

    # Check in relations (GSE series link)
    relations = summary.get("relations", [])
    for rel in relations:
        target_object = rel.get("targetobject", "")
        if target_object.startswith("GSE"):
            return target_object

    return None

def get_gse_details(gse_id):
    """
    Get detailed information for a GSE ID using E-utilities.

    Args:
        gse_id: GSE accession (e.g., GSE12345)

    Returns:
        Dictionary with dataset details
    """
    print(f"  Fetching details for {gse_id}...")

    # Search for the GSE in GEO database
    search_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
    search_params = {
        "db": "gds",
        "term": f"{gse_id}[Accession]",
        "retmode": "json"
    }

    response = requests.get(search_url, params=search_params)
    if response.status_code != 200:
        return None

    data = response.json()
    id_list = data.get("esearchresult", {}).get("idlist", [])

    if not id_list:
        return None

    # Get summary
    summary_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
    summary_params = {
        "db": "gds",
        "id": id_list[0],
        "retmode": "json"
    }

    time.sleep(0.4)  # Rate limiting
    response = requests.get(summary_url, params=summary_params)

    if response.status_code != 200:
        return None

    data = response.json()
    result = data.get("result", {}).get(id_list[0], {})

    return {
        "gse_id": gse_id,
        "title": result.get("title", ""),
        "summary": result.get("summary", ""),
        "n_samples": result.get("n_samples", 0),
        "organism": result.get("taxon", ""),
        "platform": result.get("gpl", ""),
        "pubmed_ids": result.get("pubmedids", [])
    }

def analyze_dataset_suitability(details):
    """
    Analyze if a dataset meets our criteria.

    Args:
        details: Dataset details dictionary

    Returns:
        Tuple of (is_suitable, score, metadata)
    """
    if not details:
        return False, 0, None

    title = details.get("title", "").lower()
    summary = details.get("summary", "").lower()
    n_samples = details.get("n_samples", 0)

    text = f"{title} {summary}"

    # Check for cell line
    cell_line_found = None
    for cell in CELL_LINES:
        if cell.lower() in text:
            cell_line_found = cell
            break

    # Check for compound
    compound_found = None
    for comp in COMPOUNDS:
        if comp.lower() in text:
            compound_found = comp
            break

    # Check for control/treatment indicators
    has_control = any(term in text for term in ["control", "dmso", "vehicle", "untreated"])
    has_treatment = any(term in text for term in ["treatment", "treated", "drug", "compound"])

    # Check sample count (need at least 4: 2 control + 2 treatment)
    sufficient_samples = n_samples >= 4

    # Calculate suitability score
    score = 0
    if cell_line_found:
        score += 3
    if compound_found:
        score += 3
    if has_control and has_treatment:
        score += 2
    if sufficient_samples:
        score += 2

    is_suitable = score >= 8  # Need cell line + compound + controls + samples

    metadata = {
        "cell_line": cell_line_found,
        "compound": compound_found,
        "has_controls": has_control and has_treatment,
        "n_samples": n_samples,
        "score": score
    }

    return is_suitable, score, metadata

def search_for_datasets():
    """
    Main search function to find suitable datasets.

    Returns:
        Dictionary with selected dataset information
    """
    print("=" * 80)
    print("GEO Dataset Search")
    print("=" * 80)

    # Try multiple search strategies
    candidates = []

    # Strategy 1: Search for specific combinations
    for cell_line in CELL_LINES[:3]:  # Try top 3 cell lines
        for compound in COMPOUNDS[:3]:  # Try top 3 compounds
            print(f"\nSearching: {cell_line} + {compound}")
            search_terms = [cell_line, compound, "gene expression"]

            gds_ids = search_geo_datasets(search_terms, retmax=10)

            if gds_ids:
                summaries = get_dataset_summary(gds_ids)

                for summary in summaries:
                    gse_id = parse_gse_from_summary(summary)
                    if gse_id and gse_id not in [c[0] for c in candidates]:
                        time.sleep(0.4)  # Rate limiting
                        details = get_gse_details(gse_id)
                        is_suitable, score, metadata = analyze_dataset_suitability(details)

                        if is_suitable:
                            print(f"  ✓ Found suitable dataset: {gse_id} (score: {score})")
                            candidates.append((gse_id, score, details, metadata))
                        else:
                            print(f"  ✗ Dataset {gse_id} not suitable (score: {score})")

                if len(candidates) >= 3:
                    break

        if len(candidates) >= 3:
            break

    # If no candidates found, try broader searches
    if not candidates:
        print("\nNo candidates found with specific searches. Trying broader search...")
        search_terms = ["cancer cell line", "small molecule", "gene expression", "treatment"]
        gds_ids = search_geo_datasets(search_terms, retmax=30)

        if gds_ids:
            summaries = get_dataset_summary(gds_ids)

            for summary in summaries[:15]:  # Check first 15
                gse_id = parse_gse_from_summary(summary)
                if gse_id and gse_id not in [c[0] for c in candidates]:
                    time.sleep(0.4)
                    details = get_gse_details(gse_id)
                    is_suitable, score, metadata = analyze_dataset_suitability(details)

                    if is_suitable:
                        print(f"  ✓ Found suitable dataset: {gse_id} (score: {score})")
                        candidates.append((gse_id, score, details, metadata))

                    if len(candidates) >= 3:
                        break

    # Select best candidate
    if not candidates:
        print("\n❌ No suitable datasets found!")
        return None

    # Sort by score
    candidates.sort(key=lambda x: x[1], reverse=True)

    print("\n" + "=" * 80)
    print("Candidate Datasets:")
    print("=" * 80)
    for gse_id, score, details, metadata in candidates[:5]:
        print(f"\n{gse_id} (Score: {score})")
        print(f"  Title: {details.get('title', 'N/A')[:80]}...")
        print(f"  Cell line: {metadata.get('cell_line', 'N/A')}")
        print(f"  Compound: {metadata.get('compound', 'N/A')}")
        print(f"  Samples: {metadata.get('n_samples', 'N/A')}")
        print(f"  Platform: {details.get('platform', 'N/A')}")

    # Select top candidate
    selected = candidates[0]
    gse_id, score, details, metadata = selected

    print("\n" + "=" * 80)
    print(f"Selected Dataset: {gse_id}")
    print("=" * 80)

    # Prepare output
    dataset_info = {
        "gse_id": gse_id,
        "compound": metadata.get("compound", "Unknown"),
        "cell_line": metadata.get("cell_line", "Unknown"),
        "platform": details.get("platform", "Unknown"),
        "description": details.get("title", ""),
        "n_samples": details.get("n_samples", 0),
        "organism": details.get("organism", ""),
        "summary": details.get("summary", "")[:500],  # Truncate long summaries
        "pubmed_ids": details.get("pubmed_ids", []),
        "search_score": score
    }

    return dataset_info

def main():
    """Main execution function."""
    # Search for datasets
    dataset_info = search_for_datasets()

    if dataset_info:
        # Save to JSON
        output_path = Path("/app/sandbox/session_20251217_123457_77eda5efa279/workflow/dataset_info.json")
        with open(output_path, 'w') as f:
            json.dump(dataset_info, f, indent=2)

        print(f"\n✓ Dataset information saved to: {output_path}")
        print(f"\nSelected Dataset Details:")
        print(json.dumps(dataset_info, indent=2))

        return 0
    else:
        print("\n❌ Failed to find suitable dataset!")
        return 1

if __name__ == "__main__":
    exit(main())