tensorview/state.py

"""View state serialization and deserialization for saving/loading plot configurations."""

import json
from typing import Dict, Any, Optional, List
from dataclasses import dataclass, asdict
from datetime import datetime
import orjson
from pydantic import BaseModel, Field, validator


class PlotConfig(BaseModel):
    """Plot configuration schema."""
    plot_type: str = Field(..., description="Type of plot (1d, 2d, map)")
    x_dim: Optional[str] = Field(None, description="X-axis dimension")
    y_dim: Optional[str] = Field(None, description="Y-axis dimension")
    projection: str = Field("PlateCarree", description="Map projection")
    colormap: str = Field("viridis", description="Colormap name")
    vmin: Optional[float] = Field(None, description="Color scale minimum")
    vmax: Optional[float] = Field(None, description="Color scale maximum")
    levels: Optional[List[float]] = Field(None, description="Contour levels")
    style: Dict[str, Any] = Field(default_factory=dict, description="Additional style parameters")
    
    @validator('plot_type')
    def validate_plot_type(cls, v):
        allowed = ['1d', '2d', 'map', 'contour']
        if v not in allowed:
            raise ValueError(f"plot_type must be one of {allowed}")
        return v


class DataConfig(BaseModel):
    """Data configuration schema."""
    uri: str = Field(..., description="Data source URI")
    engine: Optional[str] = Field(None, description="Data engine")
    variable_a: str = Field(..., description="Primary variable")
    variable_b: Optional[str] = Field(None, description="Secondary variable for operations")
    operation: str = Field("none", description="Operation between variables")
    selections: Dict[str, Any] = Field(default_factory=dict, description="Dimension selections")
    
    @validator('operation')
    def validate_operation(cls, v):
        allowed = ['none', 'sum', 'avg', 'diff']
        if v not in allowed:
            raise ValueError(f"operation must be one of {allowed}")
        return v


class ViewState(BaseModel):
    """Complete view state schema."""
    version: str = Field("1.0", description="State schema version")
    created: str = Field(default_factory=lambda: datetime.utcnow().isoformat(), description="Creation timestamp")
    title: str = Field("", description="User-defined title")
    description: str = Field("", description="User-defined description")
    data_config: DataConfig = Field(..., description="Data configuration")
    plot_config: PlotConfig = Field(..., description="Plot configuration")
    animation: Optional[Dict[str, Any]] = Field(None, description="Animation settings")
    exports: List[str] = Field(default_factory=list, description="Export history")
    
    class Config:
        extra = "allow"  # Allow additional fields for extensibility


def create_view_state(data_config: Dict[str, Any], plot_config: Dict[str, Any],
                     title: str = "", description: str = "",
                     animation: Optional[Dict[str, Any]] = None) -> ViewState:
    """
    Create a new view state object.
    
    Args:
        data_config: Data configuration dictionary
        plot_config: Plot configuration dictionary
        title: Optional title
        description: Optional description
        animation: Optional animation settings
        
    Returns:
        ViewState object
    """
    data_cfg = DataConfig(**data_config)
    plot_cfg = PlotConfig(**plot_config)
    
    return ViewState(
        title=title,
        description=description,
        data_config=data_cfg,
        plot_config=plot_cfg,
        animation=animation
    )


def dump_state(state: ViewState) -> str:
    """
    Serialize a view state to JSON string.
    
    Args:
        state: ViewState object
        
    Returns:
        JSON string
    """
    # Use orjson for better performance and datetime handling
    return orjson.dumps(state.dict(), option=orjson.OPT_INDENT_2).decode('utf-8')


def load_state(state_json: str) -> ViewState:
    """
    Deserialize a view state from JSON string.
    
    Args:
        state_json: JSON string
        
    Returns:
        ViewState object
    """
    try:
        data = orjson.loads(state_json)
        return ViewState(**data)
    except Exception as e:
        raise ValueError(f"Failed to parse view state: {str(e)}")


def save_state_file(state: ViewState, filepath: str) -> str:
    """
    Save view state to a file.
    
    Args:
        state: ViewState object
        filepath: Output file path
        
    Returns:
        File path
    """
    state_json = dump_state(state)
    
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(state_json)
    
    return filepath


def load_state_file(filepath: str) -> ViewState:
    """
    Load view state from a file.
    
    Args:
        filepath: Input file path
        
    Returns:
        ViewState object
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        state_json = f.read()
    
    return load_state(state_json)


def merge_states(base_state: ViewState, updates: Dict[str, Any]) -> ViewState:
    """
    Merge updates into a base state.
    
    Args:
        base_state: Base ViewState object
        updates: Dictionary of updates
        
    Returns:
        New ViewState object with updates applied
    """
    # Convert to dict, apply updates, and create new state
    state_dict = base_state.dict()
    
    # Deep merge for nested dictionaries
    def deep_merge(base_dict, update_dict):
        for key, value in update_dict.items():
            if key in base_dict and isinstance(base_dict[key], dict) and isinstance(value, dict):
                deep_merge(base_dict[key], value)
            else:
                base_dict[key] = value
    
    deep_merge(state_dict, updates)
    return ViewState(**state_dict)


def validate_state_compatibility(state: ViewState) -> List[str]:
    """
    Check state compatibility and return any warnings.
    
    Args:
        state: ViewState object
        
    Returns:
        List of warning messages
    """
    warnings = []
    
    # Check version compatibility
    current_version = "1.0"
    if state.version != current_version:
        warnings.append(f"State version {state.version} may not be fully compatible with current version {current_version}")
    
    # Check plot type and dimension compatibility
    plot_config = state.plot_config
    data_config = state.data_config
    
    if plot_config.plot_type == "map":
        if not plot_config.x_dim or not plot_config.y_dim:
            warnings.append("Map plots require both x_dim and y_dim to be specified")
    
    elif plot_config.plot_type == "2d":
        if not plot_config.x_dim or not plot_config.y_dim:
            warnings.append("2D plots require both x_dim and y_dim to be specified")
    
    elif plot_config.plot_type == "1d":
        if not plot_config.x_dim:
            warnings.append("1D plots require x_dim to be specified")
    
    # Check operation compatibility
    if data_config.operation != "none" and not data_config.variable_b:
        warnings.append(f"Operation '{data_config.operation}' requires variable_b to be specified")
    
    return warnings


def create_default_state(uri: str, variable: str) -> ViewState:
    """
    Create a default view state for a given data source and variable.
    
    Args:
        uri: Data source URI
        variable: Variable name
        
    Returns:
        Default ViewState object
    """
    data_config = {
        'uri': uri,
        'variable_a': variable,
        'operation': 'none',
        'selections': {}
    }
    
    plot_config = {
        'plot_type': '2d',
        'colormap': 'viridis',
        'style': {
            'colorbar': True,
            'grid': True
        }
    }
    
    return create_view_state(data_config, plot_config, title=f"View of {variable}")


def export_state_summary(state: ViewState) -> Dict[str, Any]:
    """
    Create a human-readable summary of a view state.
    
    Args:
        state: ViewState object
        
    Returns:
        Summary dictionary
    """
    summary = {
        'title': state.title or "Untitled View",
        'created': state.created,
        'data_source': state.data_config.uri,
        'primary_variable': state.data_config.variable_a,
        'plot_type': state.plot_config.plot_type,
        'has_secondary_variable': state.data_config.variable_b is not None,
        'operation': state.data_config.operation,
        'colormap': state.plot_config.colormap,
        'has_animation': state.animation is not None,
        'export_count': len(state.exports)
    }
    
    # Add dimension info
    selections = state.data_config.selections
    if selections:
        summary['fixed_dimensions'] = list(selections.keys())
        summary['selection_count'] = len(selections)
    
    # Add plot-specific info
    if state.plot_config.plot_type == "map":
        summary['projection'] = state.plot_config.projection
    
    return summary