"""
Correlation analysis schemas.
"""
from typing import List, Optional

from pydantic import BaseModel, Field


class CorrelationMatrixRequest(BaseModel):
    stock_codes: List[str] = Field(..., description="종목 코드 리스트")
    period_days: int = Field(60, description="분석 기간 (일)", ge=7, le=365)


class HighCorrelationPair(BaseModel):
    stock_a: str
    stock_b: str
    correlation: float = Field(..., description="상관계수 (-1 ~ 1)")


class CorrelationMatrixResponse(BaseModel):
    stock_codes: List[str]
    matrix: List[List[Optional[float]]] = Field(..., description="상관 행렬 (NxN)")
    high_correlation_pairs: List[HighCorrelationPair] = Field(
        default_factory=list, description="높은 상관관계 종목 쌍 (|r| > 0.7)"
    )


class DiversificationResponse(BaseModel):
    portfolio_id: int
    diversification_score: float = Field(
        ..., description="분산 효과 점수 (0=집중, 1=완벽 분산)", ge=0, le=1
    )
    stock_count: int
    high_correlation_pairs: List[HighCorrelationPair] = Field(default_factory=list)