from tqdm import tqdm
import torch
import copy
import numpy as np
import re
import torch.nn.functional as F


__all__ = [
    'test', 'train'
]


def test(model, test_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for test_data, i in test_loader:
            images, labels = test_data.to(device), i.to(device)
            images = images.view(images.size(0), -1)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    torch.cuda.empty_cache()
    acc = 100 * correct / total
    return acc


def train(model, train_loader, test_loader, device, criterion, optimizer, num_epochs):
    for epoch in range(num_epochs):
        model.train()
        for i, train_data in enumerate(tqdm(train_loader, desc=f'Epoch {epoch + 1}/{num_epochs}', unit='batch')):
            inputs, labels = train_data[0].to(device), train_data[1].to(device)
            inputs = inputs.view(inputs.size(0), -1)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
        torch.cuda.empty_cache()
    testacc = test(model, test_loader, device)
    print(testacc)
    return testacc



def ternarize(fp_weight, keyword, epsilon):
    ternary_weight = copy.deepcopy(fp_weight)
    for name, value in fp_weight.items():
        for kw in keyword:
            if re.match(f"^{kw}(\\.\\w+)?$", name):
                ternary_weight[name] = torch.clamp(torch.round(value / epsilon), min=-1, max=1)
    return ternary_weight
