"""
A qualified data loader should support following features
- successfully run
- len(test) == len(test_ids) == submission length
- len(train) == len(y)

Please make sure the stdout is rich enough to support informative feedback
"""

import pickle
from copy import deepcopy

import numpy as np
import pandas as pd
from feature import feat_eng
from load_data import load_data

X, y, X_test, test_ids = load_data()
print(f"X.shape: {X.shape}")
print(f"y.shape: {y.shape}" if not isinstance(y, list) else f"y(list)'s length: {len(y)}")
print(f"X_test.shape: {X_test.shape}")
print(f"test_ids length: {len(test_ids)}")
X_loaded = deepcopy(X)
y_loaded = deepcopy(y)
X_test_loaded = deepcopy(X_test)
X, y, X_test = feat_eng(X, y, X_test)


def get_length(data):
    return len(data) if isinstance(data, list) else data.shape[0]


def get_width(data):
    return 1 if isinstance(data, list) else data.shape[1:]


def get_column_list(data):
    return data.columns.tolist() if isinstance(data, pd.DataFrame) else None


assert X is not None, "The feature engineering function returned None for X."
assert y is not None, "The feature engineering function returned None for y."
assert X_test is not None, "The feature engineering function returned None for X_test."

assert get_length(X_test) == get_length(
    test_ids
), f"Mismatch in length of test images and test IDs: X_test ({get_length(X_test)}) and test_ids ({get_length(test_ids)})"
assert get_length(X) == get_length(
    y
), f"Mismatch in length of training images and labels: X ({get_length(X)}) and y ({get_length(y)})"

assert get_length(X) != 0, f"Training data is empty."
assert get_length(y) != 0, f"Training labels are empty."
assert get_length(X_test) != 0, f"Test data is empty."

assert get_width(X) == get_width(
    X_test
), "Mismatch in width of training and test data. Width means the number of features."

if isinstance(X, pd.DataFrame) and isinstance(X_test, pd.DataFrame):
    assert get_column_list(X) == get_column_list(X_test), "Mismatch in column names of training and test data."

if isinstance(X, pd.DataFrame):
    X_dtypes_unique_sorted = sorted([str(dt) for dt in X.dtypes.unique()])
    X_loaded_dtypes_unique_sorted = sorted([str(dt) for dt in X_loaded.dtypes.unique()])
    assert (
        len(X_loaded_dtypes_unique_sorted) == 1
        and X_loaded_dtypes_unique_sorted[0] in {np.float64, np.float32}
    ) or (
        X_dtypes_unique_sorted == X_loaded_dtypes_unique_sorted
    ), f"feature engineering has produced new data types which is not allowed, data loader data types are {X_loaded_dtypes_unique_sorted} and feature engineering data types are {X_dtypes_unique_sorted}"

print(
    "Feature Engineering test passed successfully. All checks including length, width, and data types have been validated."
)
