backend/neo4j/graphql_schema.py

"""
GraphQL Schema for Cancer Data
"""

import strawberry
from typing import List, Optional
from .db_manager import DatabaseManager
import logging

logger = logging.getLogger(__name__)


@strawberry.type
class Gene:
    gene_id: str
    symbol: str
    name: Optional[str] = None
    chromosome: Optional[str] = None
    start_position: Optional[int] = None
    end_position: Optional[int] = None
    gene_type: Optional[str] = None


@strawberry.type
class Mutation:
    mutation_id: str
    chromosome: str
    position: int
    reference: str
    alternate: str
    consequence: Optional[str] = None
    variant_type: Optional[str] = None
    quality: Optional[float] = None


@strawberry.type
class Patient:
    patient_id: str
    project_id: str
    age: Optional[int] = None
    gender: Optional[str] = None
    race: Optional[str] = None
    vital_status: Optional[str] = None


@strawberry.type
class CancerType:
    cancer_type_id: str
    name: str
    tissue: Optional[str] = None
    disease_type: Optional[str] = None


@strawberry.type
class MutationFrequency:
    mutation_id: str
    patients_with_mutation: int
    total_patients: int
    frequency: float


@strawberry.type
class CancerStatistics:
    cancer_type: str
    total_patients: int
    total_mutations: int
    avg_mutations_per_patient: float


@strawberry.type
class Query:
    @strawberry.field
    def gene(self, symbol: str) -> Optional[Gene]:
        """Get gene by symbol"""
        db = DatabaseManager()
        from .db_manager import GeneRepository
        repo = GeneRepository(db)
        gene_data = repo.get_gene_by_symbol(symbol)
        db.close()
        
        if gene_data:
            return Gene(**gene_data)
        return None
    
    @strawberry.field
    def genes(self, limit: int = 100) -> List[Gene]:
        """Get all genes"""
        db = DatabaseManager()
        query = "MATCH (g:Gene) RETURN g LIMIT $limit"
        results = db.execute_query(query, {'limit': limit})
        db.close()
        
        return [Gene(**r['g']) for r in results]
    
    @strawberry.field
    def mutations(
        self,
        gene: Optional[str] = None,
        chromosome: Optional[str] = None,
        limit: int = 100
    ) -> List[Mutation]:
        """Get mutations, optionally filtered by gene or chromosome"""
        db = DatabaseManager()
        
        if gene:
            query = """
            MATCH (g:Gene {symbol: $gene})<-[:AFFECTS]-(m:Mutation)
            RETURN m
            LIMIT $limit
            """
            params = {'gene': gene, 'limit': limit}
        elif chromosome:
            query = """
            MATCH (m:Mutation {chromosome: $chromosome})
            RETURN m
            LIMIT $limit
            """
            params = {'chromosome': chromosome, 'limit': limit}
        else:
            query = "MATCH (m:Mutation) RETURN m LIMIT $limit"
            params = {'limit': limit}
        
        results = db.execute_query(query, params)
        db.close()
        
        return [Mutation(**r['m']) for r in results]
    
    @strawberry.field
    def patients(
        self,
        project_id: Optional[str] = None,
        cancer_type: Optional[str] = None,
        limit: int = 100
    ) -> List[Patient]:
        """Get patients, optionally filtered"""
        db = DatabaseManager()
        
        if project_id:
            query = """
            MATCH (p:Patient {project_id: $project_id})
            RETURN p
            LIMIT $limit
            """
            params = {'project_id': project_id, 'limit': limit}
        elif cancer_type:
            query = """
            MATCH (p:Patient)-[:DIAGNOSED_WITH]->(c:CancerType {cancer_type_id: $cancer_type})
            RETURN p
            LIMIT $limit
            """
            params = {'cancer_type': cancer_type, 'limit': limit}
        else:
            query = "MATCH (p:Patient) RETURN p LIMIT $limit"
            params = {'limit': limit}
        
        results = db.execute_query(query, params)
        db.close()
        
        return [Patient(**r['p']) for r in results]
    
    @strawberry.field
    def cancer_types(self) -> List[CancerType]:
        """Get all cancer types"""
        db = DatabaseManager()
        query = "MATCH (c:CancerType) RETURN c"
        results = db.execute_query(query)
        db.close()
        
        return [CancerType(**r['c']) for r in results]
    
    @strawberry.field
    def mutation_frequency(self, mutation_id: str) -> Optional[MutationFrequency]:
        """Get frequency of a mutation across all patients"""
        db = DatabaseManager()
        from .db_manager import MutationRepository
        repo = MutationRepository(db)
        freq_data = repo.get_mutation_frequency(mutation_id)
        db.close()
        
        if freq_data:
            return MutationFrequency(**freq_data)
        return None
    
    @strawberry.field
    def cancer_statistics(self, cancer_type_id: str) -> Optional[CancerStatistics]:
        """Get statistics for a cancer type"""
        db = DatabaseManager()
        from .db_manager import CancerTypeRepository
        repo = CancerTypeRepository(db)
        stats = repo.get_statistics(cancer_type_id)
        db.close()
        
        if stats:
            return CancerStatistics(**stats)
        return None


schema = strawberry.Schema(query=Query)