import numpy as np
from sklearn.ensemble import RandomForestClassifier
import dataclasses
import matplotlib.pyplot as plt
from sklearn import metrics
import json


@dataclasses.dataclass
class Config:
    train_images_path: str = 'data/processed/train-images-idx3-ubyte'
    train_labels_path: str = 'data/processed/train-labels-idx1-ubyte'
    train_sample_size: int = 60000
    
    test_images_path: str  = 'data/processed/t10k-images-idx3-ubyte'
    test_labels_path: str  = 'data/processed/t10k-labels-idx1-ubyte'
    test_sample_size: int  = 10000
    
    with_unsupervised_learning_trick: bool = True
    seed: int = 42
    train_n_unknowns: int = 6000
    test_n_unknowns: int = 1000
    
    random_forest_kwargs: dict = dataclasses.field(
        default_factory=lambda:{
            'n_estimators': 500,
            'max_depth': 13,
            'random_state': 42,
            'n_jobs': -1,
        }
    )

config = Config()

def load_images(images_path):
    with open(images_path, 'rb') as f:
        magic_number = int.from_bytes(f.read(4), 'big')
        n_images = int.from_bytes(f.read(4), 'big')
        n_columns = int.from_bytes(f.read(4), 'big')
        n_rows = int.from_bytes(f.read(4), 'big')
        images = np.frombuffer(f.read(), dtype=np.uint8)
    images = images.reshape(n_images, n_columns * n_rows)
    return magic_number, n_images, n_columns, n_rows, images


def load_labels(labels_path):
    with open(labels_path, 'rb') as f:
        magic_number = int.from_bytes(f.read(4), 'big')
        n_items = int.from_bytes(f.read(4), 'big')
        labels = np.frombuffer(f.read(), dtype=np.uint8)
    return magic_number, n_items, labels


def load_train(config):
    _, _, _, _, images = load_images(
        images_path=config.train_images_path
    )
    _, _, labels = load_labels(
        labels_path=config.train_labels_path,
    )
    images = images[0:config.train_sample_size,:]
    labels = labels[0:config.train_sample_size]
    return images, labels


def load_test(config):
    _, _, _, _, images = load_images(
        images_path=config.test_images_path
    )
    _, _, labels = load_labels(
        labels_path=config.test_labels_path,
    )
    images = images[0:config.test_sample_size,:]
    labels = labels[0:config.test_sample_size]
    return images, labels

def unsupervised_learning_trick(X, y, n_unknowns):
    
    # Create the simulated observations.
    features = []
    for feature_idx in range(X_train.shape[1]):
        
        # Randomly select the value for variable from all the observed values of variable.
        random_idxs = np.random.choice(X_train.shape[0], size=n_unknowns)
        feature = X_train[random_idxs, feature_idx]
        
        # Save.
        features.append(feature)
        
    # Combine. Unknowns are assigned label -1.
    X_unknown = np.column_stack(features)
    y_unknown = np.repeat(-1, n_unknowns)
    
    # Concatenate with original data.
    X = np.r_[X, X_unknown]
    y = np.r_[y, y_unknown]
    
    return X, y


def evaluate(y_true, y_pred):
    print(metrics.classification_report(y_true, y_pred))
    fig, ax = plt.subplots(figsize=(12, 12))
    metrics.ConfusionMatrixDisplay.from_predictions(y_true, y_pred, ax=ax, normalize='true')
    plt.show()
    
    
def eyeball(X, y, y_hat, y_prob, n_max_eyeballs=10):
    for row_idx in range(min(X.shape[0], n_max_eyeballs)):
        print('True Label', y[row_idx])
        print('Predicted Label', y_hat[row_idx])
        print('Predicted Probability', np.max(y_prob[row_idx]))
        plt.figure()
        plt.imshow(X[row_idx, :].reshape(28, 28))
        plt.show()

# Set a seed for reproducibility.
np.random.seed(config.seed)

X_train, y_train = load_train(config)
if config.with_unsupervised_learning_trick:
    X_train, y_train = unsupervised_learning_trick(X_train, y_train, n_unknowns=config.train_n_unknowns)
rf = RandomForestClassifier(**config.random_forest_kwargs)
rf.fit(X_train, y_train)

y_hat_train = rf.predict(X_train)
evaluate(y_train, y_hat_train)

X_test, y_test = load_test(config)
if config.with_unsupervised_learning_trick:
    X_test, y_test = unsupervised_learning_trick(X_test, y_test, n_unknowns=config.test_n_unknowns)
y_hat_test = rf.predict(X_test)
y_prob_test = rf.predict_proba(X_test)
evaluate(y_test, y_hat_test)

not_unknown_mask = y_test != -1
print("Total", sum(not_unknown_mask))
print()
eyeball(
    X=X_test[not_unknown_mask],
    y=y_test[not_unknown_mask],
    y_hat=y_hat_test[not_unknown_mask],
    y_prob=y_prob_test[not_unknown_mask]
)

unknown_mask = y_hat_test == -1
print("Total", sum(unknown_mask))
print()
eyeball(
    X=X_test[unknown_mask],
    y=y_test[unknown_mask],
    y_hat=y_hat_test[unknown_mask],
    y_prob=y_prob_test[unknown_mask],
)