# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_record_validation.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/03_record_validation.ipynb 2
import pandas as pd
from pandas import DataFrame as PandasDataFrame
from typing import *
from archetypon.base_model import BaseModel,GenericModel,DataFrame
from pydantic import parse_obj_as,ValidationError,validator
from pydantic.utils import update_not_none

# %% ../nbs/03_record_validation.ipynb 3
def parse_dataframe_rows_as(
    model:Type[BaseModel],
    df:PandasDataFrame
)->PandasDataFrame:
    """Uses .parse_obj() method of Pydantic's `BaseModel` to validate rows of a dataframe."""
    
    # convert dataframe to a series of dictionaries
    # drop NaN values because pydantic doesn't consider them to be 'None', and that interferes with validation
    series_of_dicts = df.apply(
        lambda row: row.dropna().to_dict(),
        axis=1
    )
    
    #convert the series of dicts to a series of parsed models
    errors = []
    def parse_row(row):
        try:
            validated = parse_obj_as(model,row)
            return validated
        except ValidationError as e:
            raw_errors = e.raw_errors
            
            errors.extend(raw_errors)
            
            return BaseModel()
        return validated
    
    series_of_models = series_of_dicts.apply(
        lambda x: parse_row(x),
    )
    if len(errors)>0:
        raise ValidationError(errors=errors,model=model)
    # from model back to dictionaries and then back to Series
    validated = series_of_models.apply(
        lambda x: x.dict(),
    ).apply(pd.Series)
    
    return validated

# %% ../nbs/03_record_validation.ipynb 9
class TypedRecordFrame(DataFrame):
    row_model: Optional[Type[BaseModel]] = None
    alias_as_column_names: bool = False
        
    @classmethod
    def __get_validators__(cls):
        for v in super().__get_validators__():
            yield v
        yield cls.validate_rows
        yield cls.validate_column_names
    
    @classmethod
    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
        row_model=None
        if cls.row_model:
            row_model = cls.row_model.schema()
        update_not_none(
            field_schema,
            row_model=row_model
        )

    @classmethod
    def validate_rows(cls, df):
        
        if cls.row_model:
            validated = parse_dataframe_rows_as(cls.row_model,df)
                            
            return validated
        return df
    
    @classmethod
    def validate_column_names(cls,df):
        if cls.alias_as_column_names==True:
            # create a dictionary mapping field names to aliases
            field_name_to_alias = {
                field_name: cls.row_model.__fields__[field_name].alias
                for field_name in cls.row_model.__fields__
            }

            df.rename(columns=field_name_to_alias,inplace=True)
        return df
            

# %% ../nbs/03_record_validation.ipynb 10
class RecordModelFrameMeta(type):
    def __getitem__(self, constraint):
        return type('RecordFrame', (TypedRecordFrame,), {'row_model': constraint})


# %% ../nbs/03_record_validation.ipynb 11
class RecordFrame(DataFrame, metaclass=RecordModelFrameMeta):
    pass
