from dataclasses import dataclass, field
from typing import Any, Union
import pandas as pd


Array = pd.Series
DataType = pd.DataFrame


@dataclass(frozen=True)
class Schema:
    name: str
    dtype: type
    not_null: bool = field(default=True)


class Dataset:
    def __init__(self, data: DataType, schema: dict[str, type] = None) -> None:
        self.data = data
        self.schema = schema if schema else dict() 
    
    def __repr__(self):
        return f"{self.data}"
    
    @classmethod
    def from_csv(cls, path: str) -> 'Dataset':
        dataset: DataType = pd.read_csv(path)
        return cls(dataset)
    
    @classmethod
    def from_dict(cls, data: dict[str, list]) -> 'Dataset':
        dataset: DataType = pd.DataFrame(data)
        return cls(dataset)

    @classmethod
    def empty(cls, schema: dict[str, type]) -> 'Dataset':
        # dataset: DataType = pd.DataFrame(columns=list(schema.keys()))
        dataset: DataType = pd.DataFrame({name: pd.Series(dtype=dtype) for name, dtype in schema.items()})
        return cls(dataset, schema)
    
    def to_csv(self, path: str) -> None:
        self.data.to_csv(path, header=True, index=False)
    
    def to_parquet(self, path: str) -> None:
        self.data.to_parquet(path, compression="brotli", header=True, index=False)
    
    def col(self, name: str) -> Array:
        return self.data[name]

    def append(self, data: dict[str, Any]) -> None:
        new_data = pd.DataFrame({name: pd.Series(data=data[name], dtype=dtype) for name, dtype in self.schema.items()})
        if len(self.data) > 0:
            self.data = pd.concat([self.data, new_data], ignore_index=True)
        else:
            self.data = new_data

    def normalize(self, cols: list[str]) -> 'Dataset':
        new_data = self.data.copy()

        for col in cols:
            x = new_data[col]
            x_min = x.min()
            x_max = x.max()
            
            new_data[col] = (x - x_min)/(x_max - x_min)

        return Dataset.from_dict(data=new_data)

    def standardize(self, cols: list[str]) -> 'Dataset':
        new_data = self.data.copy()

        for col in cols:
            x = new_data[col]
            u = x.mean()
            s = x.std()
            
            new_data[col] = (x - u)/s

        return Dataset.from_dict(data=new_data)

