import multiprocessing as mp
from typing import Tuple

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, random_split
from tqdm.auto import tqdm
TRAIN_FILES = "/kaggle/input/digit-recognizer/train.csv"
TRAIN_LEN = 0.8
IMAGE_SIZE = (28, 28)
BATCH_SIZE = 64
EPOCHS = 10
LEARNING_RATE = 1e-4
"""
## Data
"""
class MNISTData(Dataset):
    def __init__(self, path: str):
        df = pd.read_csv(path)
        self.x = df.iloc[:, 1:].values / 255.0
        self.y = df.iloc[:, 0].values / 10.0 - 0.5
        
    def __len__(self):
        return len(self.y)
    
    def __getitem__(self, i: int) -> Tuple[np.array, np.array]:
        return self.x[i], self.y[i]
mnist = MNISTData(TRAIN_FILES)
train_len = int(len(mnist) * TRAIN_LEN)
valid_len = len(mnist) - train_len
train_ds, valid_ds = random_split(mnist, [train_len, valid_len])

train_dl = DataLoader(
    train_ds,
    BATCH_SIZE,
    num_workers=mp.cpu_count(),
    shuffle=True,
    drop_last=True,
)

valid_dl = DataLoader(
    valid_ds,
    BATCH_SIZE,
    num_workers=mp.cpu_count(),
    shuffle=False,
    drop_last=False,
)
x, y = next(iter(train_dl))
x.shape, y.shape, x.dtype, y.dtype
"""
## Model
"""
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(784, 100)
        self.linear2 = nn.Linear(100, 1)
        self.bn1 = nn.BatchNorm1d(100)
        
    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
        breakpoint()
        x = self.bn1(nn.ReLU()(self.linear1(x)))
        return self.linear2(x)
model = Model()
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), LEARNING_RATE)
for _ in range(EPOCHS):
    pbar = tqdm(train_dl)
    for x, y in pbar:
        y_pred = model(x.float())
        loss = loss_fn(y_pred.squeeze(), y.float().squeeze())
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        pbar.set_description(f"Loss is {loss.item():.4f}")
#         print(f"\r", end="")
y_preds = []
ys = []
model = model.eval()
with torch.inference_mode():
    for x, y in tqdm(valid_dl):
        y_preds.extend(model(x.float()))
        ys.extend(y.float())
        
y_preds = 10.0 * (torch.stack(y_preds) + 0.5)
ys = 10.0 * (torch.stack(ys) + 0.5)
(y_preds.round().squeeze() == ys).float().mean()