database/models.py

"""
Database models for the AI Knowledge Distillation Platform
"""

from dataclasses import dataclass
from typing import Optional, Dict, Any, List
from datetime import datetime
import json

@dataclass
class TokenModel:
    """Model for HF token storage"""
    id: Optional[int] = None
    name: str = ""
    token_type: str = "read"
    encrypted_token: str = ""
    is_default: bool = False
    description: str = ""
    created_at: Optional[datetime] = None
    last_used: Optional[datetime] = None
    usage_count: int = 0
    is_active: bool = True
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'id': self.id,
            'name': self.name,
            'token_type': self.token_type,
            'encrypted_token': self.encrypted_token,
            'is_default': self.is_default,
            'description': self.description,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'last_used': self.last_used.isoformat() if self.last_used else None,
            'usage_count': self.usage_count,
            'is_active': self.is_active
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'TokenModel':
        """Create from dictionary"""
        return cls(
            id=data.get('id'),
            name=data.get('name', ''),
            token_type=data.get('token_type', 'read'),
            encrypted_token=data.get('encrypted_token', ''),
            is_default=data.get('is_default', False),
            description=data.get('description', ''),
            created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at') else None,
            last_used=datetime.fromisoformat(data['last_used']) if data.get('last_used') else None,
            usage_count=data.get('usage_count', 0),
            is_active=data.get('is_active', True)
        )

@dataclass
class TrainingSessionModel:
    """Model for training session data"""
    id: Optional[int] = None
    session_id: str = ""
    teacher_model: str = ""
    student_model: str = ""
    dataset_name: Optional[str] = None
    training_type: str = "knowledge_distillation"
    status: str = "initialized"
    progress: float = 0.0
    current_step: int = 0
    total_steps: Optional[int] = None
    current_loss: Optional[float] = None
    best_loss: Optional[float] = None
    learning_rate: Optional[float] = None
    batch_size: Optional[int] = None
    temperature: Optional[float] = None
    alpha: Optional[float] = None
    created_at: Optional[datetime] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error_message: Optional[str] = None
    config: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'id': self.id,
            'session_id': self.session_id,
            'teacher_model': self.teacher_model,
            'student_model': self.student_model,
            'dataset_name': self.dataset_name,
            'training_type': self.training_type,
            'status': self.status,
            'progress': self.progress,
            'current_step': self.current_step,
            'total_steps': self.total_steps,
            'current_loss': self.current_loss,
            'best_loss': self.best_loss,
            'learning_rate': self.learning_rate,
            'batch_size': self.batch_size,
            'temperature': self.temperature,
            'alpha': self.alpha,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None,
            'error_message': self.error_message,
            'config': self.config
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'TrainingSessionModel':
        """Create from dictionary"""
        return cls(
            id=data.get('id'),
            session_id=data.get('session_id', ''),
            teacher_model=data.get('teacher_model', ''),
            student_model=data.get('student_model', ''),
            dataset_name=data.get('dataset_name'),
            training_type=data.get('training_type', 'knowledge_distillation'),
            status=data.get('status', 'initialized'),
            progress=data.get('progress', 0.0),
            current_step=data.get('current_step', 0),
            total_steps=data.get('total_steps'),
            current_loss=data.get('current_loss'),
            best_loss=data.get('best_loss'),
            learning_rate=data.get('learning_rate'),
            batch_size=data.get('batch_size'),
            temperature=data.get('temperature'),
            alpha=data.get('alpha'),
            created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at') else None,
            started_at=datetime.fromisoformat(data['started_at']) if data.get('started_at') else None,
            completed_at=datetime.fromisoformat(data['completed_at']) if data.get('completed_at') else None,
            error_message=data.get('error_message'),
            config=data.get('config')
        )
    
    def get_config_json(self) -> str:
        """Get config as JSON string"""
        return json.dumps(self.config) if self.config else ""
    
    def set_config_from_json(self, config_json: str):
        """Set config from JSON string"""
        try:
            self.config = json.loads(config_json) if config_json else None
        except json.JSONDecodeError:
            self.config = None

@dataclass
class PerformanceMetricModel:
    """Model for performance metrics"""
    id: Optional[int] = None
    timestamp: Optional[datetime] = None
    metric_type: str = "system"  # system, model, training
    metric_name: str = ""
    metric_value: float = 0.0
    unit: str = ""
    context: Optional[str] = None  # Additional context (model name, session id, etc.)
    metadata: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'id': self.id,
            'timestamp': self.timestamp.isoformat() if self.timestamp else None,
            'metric_type': self.metric_type,
            'metric_name': self.metric_name,
            'metric_value': self.metric_value,
            'unit': self.unit,
            'context': self.context,
            'metadata': self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'PerformanceMetricModel':
        """Create from dictionary"""
        return cls(
            id=data.get('id'),
            timestamp=datetime.fromisoformat(data['timestamp']) if data.get('timestamp') else None,
            metric_type=data.get('metric_type', 'system'),
            metric_name=data.get('metric_name', ''),
            metric_value=data.get('metric_value', 0.0),
            unit=data.get('unit', ''),
            context=data.get('context'),
            metadata=data.get('metadata')
        )

@dataclass
class MedicalDatasetModel:
    """Model for medical dataset information"""
    id: Optional[int] = None
    dataset_name: str = ""
    repo_id: str = ""
    description: str = ""
    size_gb: float = 0.0
    num_samples: int = 0
    modalities: List[str] = None
    specialties: List[str] = None
    languages: List[str] = None
    last_accessed: Optional[datetime] = None
    access_count: int = 0
    is_cached: bool = False
    cache_path: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    
    def __post_init__(self):
        """Initialize default values"""
        if self.modalities is None:
            self.modalities = []
        if self.specialties is None:
            self.specialties = []
        if self.languages is None:
            self.languages = []
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'id': self.id,
            'dataset_name': self.dataset_name,
            'repo_id': self.repo_id,
            'description': self.description,
            'size_gb': self.size_gb,
            'num_samples': self.num_samples,
            'modalities': self.modalities,
            'specialties': self.specialties,
            'languages': self.languages,
            'last_accessed': self.last_accessed.isoformat() if self.last_accessed else None,
            'access_count': self.access_count,
            'is_cached': self.is_cached,
            'cache_path': self.cache_path,
            'metadata': self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MedicalDatasetModel':
        """Create from dictionary"""
        return cls(
            id=data.get('id'),
            dataset_name=data.get('dataset_name', ''),
            repo_id=data.get('repo_id', ''),
            description=data.get('description', ''),
            size_gb=data.get('size_gb', 0.0),
            num_samples=data.get('num_samples', 0),
            modalities=data.get('modalities', []),
            specialties=data.get('specialties', []),
            languages=data.get('languages', []),
            last_accessed=datetime.fromisoformat(data['last_accessed']) if data.get('last_accessed') else None,
            access_count=data.get('access_count', 0),
            is_cached=data.get('is_cached', False),
            cache_path=data.get('cache_path'),
            metadata=data.get('metadata')
        )
    
    def get_modalities_string(self) -> str:
        """Get modalities as comma-separated string"""
        return ','.join(self.modalities) if self.modalities else ""
    
    def get_specialties_string(self) -> str:
        """Get specialties as comma-separated string"""
        return ','.join(self.specialties) if self.specialties else ""
    
    def get_languages_string(self) -> str:
        """Get languages as comma-separated string"""
        return ','.join(self.languages) if self.languages else ""
    
    def set_modalities_from_string(self, modalities_str: str):
        """Set modalities from comma-separated string"""
        self.modalities = [m.strip() for m in modalities_str.split(',') if m.strip()] if modalities_str else []
    
    def set_specialties_from_string(self, specialties_str: str):
        """Set specialties from comma-separated string"""
        self.specialties = [s.strip() for s in specialties_str.split(',') if s.strip()] if specialties_str else []
    
    def set_languages_from_string(self, languages_str: str):
        """Set languages from comma-separated string"""
        self.languages = [l.strip() for l in languages_str.split(',') if l.strip()] if languages_str else []

@dataclass
class DicomFileModel:
    """Model for DICOM file information"""
    id: Optional[int] = None
    file_path: str = ""
    patient_id: Optional[str] = None
    study_date: Optional[str] = None
    modality: Optional[str] = None
    file_size_mb: float = 0.0
    processed: bool = False
    processed_at: Optional[datetime] = None
    metadata: Optional[Dict[str, Any]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            'id': self.id,
            'file_path': self.file_path,
            'patient_id': self.patient_id,
            'study_date': self.study_date,
            'modality': self.modality,
            'file_size_mb': self.file_size_mb,
            'processed': self.processed,
            'processed_at': self.processed_at.isoformat() if self.processed_at else None,
            'metadata': self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'DicomFileModel':
        """Create from dictionary"""
        return cls(
            id=data.get('id'),
            file_path=data.get('file_path', ''),
            patient_id=data.get('patient_id'),
            study_date=data.get('study_date'),
            modality=data.get('modality'),
            file_size_mb=data.get('file_size_mb', 0.0),
            processed=data.get('processed', False),
            processed_at=datetime.fromisoformat(data['processed_at']) if data.get('processed_at') else None,
            metadata=data.get('metadata')
        )