from typing import Dict, List, Optional

from pydantic import BaseModel


class SegmentResult(BaseModel):
    segment: Dict[str, str]
    total: float
    seg_size: float
    naive_avg: float
    impact: Optional[float] = None
    avg_impact: Optional[float] = None


class AnalysisResult(BaseModel):
    task: str
    segments: List[SegmentResult]
    relevant_clusters: Dict[str, str]
    global_average: Optional[float] = None
    pre_total: Optional[float] = None
    post_total: Optional[float] = None
    markdown_summary: Optional[str] = None
    image_path: Optional[str] = None


class SplitAnalysisResult(BaseModel):
    task: str
    size_analysis: AnalysisResult
    average_analysis: AnalysisResult
    image_path: Optional[str] = None


class LoadDataResult(BaseModel):
    dataset_name: str
    num_rows: int
    num_columns: int
    columns: List[str]
    dims: List[str]
    total_name: str
    size_name: Optional[str] = None
