# Simple model for each beam

import numpy  as np

import pandas as pd

from   sklearn.ensemble import HistGradientBoostingRegressor

# Read data

ss     = pd.read_csv ( "SampleSubmission.csv" )

ty     = pd.read_csv ( "traffic_DLThpVol.csv" )

sub_fn = "simple dummy sub.csv"

# Configuration

n_base =        30

n_cell =         3

n_beam =        32

rs     =  123

# Dummies

def make_dummies ( n ) :

    i  = np.arange ( n, dtype = int )

    df = pd.DataFrame ( index = i )

    # Hour

    for j in range ( 24 - 1 ) :

        df [ f"h{ j }" ] = 1.0 * ( ( i % 24 ) == j )

    # Week

    for j in range ( 7 - 1 ) :

        df [ f"w{ j }" ] = 1.0 * ( ( ( i // 24 ) % 7 ) == j )

    return df

x_train = make_dummies ( len ( ty ) )

x_test  = make_dummies ( len ( ty ) + 1008 ).iloc [ len ( ty ) : ]

# Prepare sample submission

ss = ss.set_index ( "ID", drop = False )

ss [ "Target" ] = ss [ "Target" ].astype ( "float16" )

# Fit one model for each beam

for base in range ( n_base ) :

    for cell in range ( n_cell ) :

        for beam in range ( n_beam ) :

            rs     += 1

            mod_col = f"{ base }_{ cell }_{ beam }"

            model   = HistGradientBoostingRegressor ( random_state = rs, loss = "absolute_error" )

            model.fit ( x_train, ty [ mod_col ] )

            pred    = np.clip ( model.predict ( x_test ), 0, 255 )

            # Load predictions into sample submission

            for k in range ( 168 ) :

                ss.at [ "traffic_DLThpVol_test_5w-6w_"   + str (       k     ) + "_" + mod_col, "Target" ] = pred [   k     ]

                ss.at [ "traffic_DLThpVol_test_10w-11w_" + str ( 168 - k - 1 ) + "_" + mod_col, "Target" ] = pred [ - k - 1 ]

# Save submission

ss.to_csv ( sub_fn, index = False )