<gh_stars>0
# Databricks notebook source
# MAGIC %run ../../app/bootstrap

# COMMAND ----------

from pyspark.sql.dataframe import DataFrame
from datalakebundle.imports import transformation, table_overwrite, read_csv, read_table
import pyspark.sql.functions as f
import pyspark.sql.types as t

from daipeproject.bronze.visits.schema import get_schema as get_csv_schema
from daipeproject.silver.visits.schema import get_schema as get_table_schema

# COMMAND ----------

def get_unique_custom_attrs(schema: t.StructType):
    return [field.name for field in schema.fields if "custom_field" in field.metadata and field.metadata["custom_field"] is True]
  
def get_mapping_col_name(name: str):
    return name + "_column"

# COMMAND ----------

@transformation("%datasets%", display=True)
def custom_attrs_mapping(datasets):
    unique_custom_attrs = get_unique_custom_attrs(get_table_schema())
    schema = ["name"] + [get_mapping_col_name(name) for name in unique_custom_attrs]
  
    data = []
    
    for dataset in datasets:
        row = tuple([dataset.name] + [dataset.custom_attrs[attr] if attr in dataset.custom_attrs else None for attr in unique_custom_attrs])
      
        data.append(row)
  
    return spark.createDataFrame(data, schema=schema)

# COMMAND ----------

@transformation(
    custom_attrs_mapping,
    read_csv(
        "%datalake.base_path%/bronze/raw/visits", # or "%datalake.base_path%/bronze/raw/visits/*" (without name as partition_key)
        schema=get_csv_schema(),
        options=dict(header='true')
    ),
    display=True
)
@table_overwrite("daipe_demo_ga_db.visits", get_table_schema(), recreate_table=True)
def process_visits(custom_attrs_mapping_df: DataFrame, df: DataFrame):    
    table_schema = get_table_schema()
    
    custom_x_cols = [field.name for field in df.schema.fields if field.name[:7] == "custom_"]
    
    def create_custom_column(custom_attr: str, custom_x_cols: list):
        mapping_col = get_mapping_col_name(custom_attr)
        conditions = [f.when(f.col(mapping_col) == custom_x_col, f.col(custom_x_col)) for custom_x_col in custom_x_cols]
        datatype = table_schema[custom_attr].dataType
        
        return f.coalesce(*conditions).cast(datatype).alias(custom_attr)
    
    select_columns = [
        f.col("name"),
        f.to_date(f.col("date"), 'yyyy-MM-dd').alias("date"),
        f.col("visits"),
    ] + [create_custom_column(custom_attr, custom_x_cols) for custom_attr in get_unique_custom_attrs(table_schema)]
    
    return (
        df
           .join(custom_attrs_mapping_df, on="name")
           .select(*select_columns)
           .sort(f.col("name"), f.col("date").asc())
    )
