"""
A qualified ensemble implementation should:
- Successfully run
- Return predictions
- Have correct shapes for inputs and outputs
- Use validation data appropriately
"""

import numpy as np
import pandas as pd
from pathlib import Path
from sklearn.model_selection import train_test_split
import torch
import tensorflow as tf
from load_data import load_data
from feature import feat_eng
from ensemble import ensemble_workflow

X, y, test_X, test_ids = load_data()
X, y, test_X = feat_eng(X, y, test_X)
train_X, val_X, train_y, val_y = train_test_split(X, y, test_size=0.2, random_state=42)

# Print the types of train_y and val_y
print(f"train_y type: {type(train_y)}, val_y type: {type(val_y)}")

test_preds_dict = {}
val_preds_dict = {}
{% for mn in model_names %}
from {{mn}} import model_workflow as {{mn}}_workflow
val_preds_dict["{{mn}}"], test_preds_dict["{{mn}}"], _ = {{mn}}_workflow(
    X=train_X,
    y=train_y,
    val_X=val_X,
    val_y=val_y,
    test_X=test_X
)
{% endfor %}

for key in val_preds_dict.keys():
    if val_preds_dict[key] is None: 
        print(f"Model {key} validation predictions (val_preds_dict[key]) is None.")
    elif isinstance(val_preds_dict[key], list):
        print(f"Model {key} validation predictions (val_preds_dict[key]) (list type) length: {len(val_preds_dict[key])}")
    else:
        print(f"Model {key} validation predictions (val_preds_dict[key]) shape: {val_preds_dict[key].shape}")

    if test_preds_dict[key] is None: 
        print(f"Model {key} test predictions (test_preds_dict[key]) is None.")
    elif isinstance(test_preds_dict[key], list):
        print(f"Model {key} test predictions (test_preds_dict[key]) (list type) length: {len(test_preds_dict[key])}")
    else:
        print(f"Model {key} test predictions (test_preds_dict[key]) shape: {test_preds_dict[key].shape}")

# Run ensemble
final_pred = ensemble_workflow(test_preds_dict, val_preds_dict, val_y)

# Check type
pred_type = type(next(iter(test_preds_dict.values())))
assert isinstance(final_pred, pred_type), (
    f"Type mismatch: 'final_pred' is of type {type(final_pred)}, but expected {pred_type} "
)

# Check shape
if isinstance(final_pred, (list, np.ndarray, pd.DataFrame, torch.Tensor, tf.Tensor)):
    assert len(final_pred) == len(test_X), (
        f"Wrong output sample size: len(final_pred)={len(final_pred)} "
        f"vs. len(test_X)={len(test_X)}"
    )

# check scores.csv
assert Path("scores.csv").exists(), "scores.csv is not generated"
score_df = pd.read_csv("scores.csv", index_col=0)
model_set_in_scores = set(score_df.index)
for model in {{model_names}}:
    if model not in model_set_in_scores:
        print(f"\nModel {model} is not evaluated in the scores.csv.")
print("Please check scores dataframe's format, it's:")
print(score_df)

print("Ensemble test end.")
print(f"Final prediction shape: {final_pred.shape}")
