fix the checkpoint from IPython.display import clear_output
from pathlib import Path
from ray import tune, air
from ray.air import session, Checkpoint
from ray.tune.schedulers import PopulationBasedTraining
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import Dataset, DataLoader
from torchvision.models import resnet50
from tqdm import tqdm
import cv2
import json
import math
import numpy as np
import optuna
import os
import random
import ray
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms


# Set Up
#!device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')

RAY_USE_XRAY=1 

# Set a random seed for torch
seed = 13
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)  # If using multiple GPUs
np.random.seed(seed)
random.seed(seed)

# Additional configurations for deterministic behavior on CUDA devices
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

data_path = '/workspaces/local_dev/Images/Training'
validation_path = '/workspaces/local_dev/Images/Validation'
model_name = 'rotnet_resnet50v2'

# Training Parameters
num_samples = 100  # None # Number of samples used for training. If None then all images
# None  # Number of samples used for validation. If None then all images
num_val_samples = 25
# Tuple defining the input image shape (height, width)
input_shape = (224, 224)
num_classes = 4  # The possible rotation angles: 0, 90, 180, 270

# Optuna Optimization Parameters
n_trials = 5  # Number of Optuna Trials
learning_rate_options = (1e-5, 1e-1)  # Learning rate for the optimizer
batch_size_options = [8, 16]  # Batch size for training
num_epochs_options = (5, 5)  # Number of epochs to train the model for
weight_decay_options = (1e-6, 0.1)
dropout_rate_options = (0.001, 0.5)

# Create output folder if it doesn't exist
output_folder = 'models'
Path(output_folder).mkdir(parents=True, exist_ok=True)

# Define Defs

# Model definition


class RotNetModel(nn.Module):
    def __init__(self, num_classes):
        super(RotNetModel, self).__init__()
        self.resnet = resnet50(weights=True)
        self.resnet.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        return self.resnet(x)

# Error Function


def angle_difference(x, y):
    """
    Calculate minimum difference between two angles.
    """
    return 180 - abs(abs(x - y) - 180)


def angle_error(y_true, y_pred):
    """
    Calculate the mean difference between the true angles
    and the predicted angles.
    """
    diff = angle_difference(y_true, y_pred)
    return torch.mean(torch.abs(diff.float()))

# Rotation Related


def rotate(image, angle):
    """
    Rotates an OpenCV 2 / NumPy image about its center by the given angle
    (in degrees).
    """
    center = (image.shape[1] / 2, image.shape[0] / 2)
    rot_mat = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(image, rot_mat, (image.shape[1], image.shape[0]), flags=cv2.INTER_LINEAR)


def largest_rotated_rect(w, h, angle):
    """
    Given a rectangle of size wxh that has been rotated by 'angle' (in
    radians), computes the width and height of the largest possible
    axis-aligned rectangle within the rotated rectangle.

    Original JS code by 'Andri' and Magnus Hoff from Stack Overflow

    Converted to Python by Aaron Snoswell

    Source: http://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders
    """

    quadrant = int(math.floor(angle / (math.pi / 2))) & 3
    sign_alpha = angle if ((quadrant & 1) == 0) else math.pi - angle
    alpha = (sign_alpha % math.pi + math.pi) % math.pi

    bb_w = w * math.cos(alpha) + h * math.sin(alpha)
    bb_h = w * math.sin(alpha) + h * math.cos(alpha)

    gamma = math.atan2(bb_w, bb_w) if (w < h) else math.atan2(bb_w, bb_w)

    delta = math.pi - alpha - gamma

    length = h if (w < h) else w

    d = length * math.cos(alpha)
    a = d * math.sin(alpha) / math.sin(delta)

    y = a * math.cos(gamma)
    x = y * math.tan(gamma)

    return (
        bb_w - 2 * x,
        bb_h - 2 * y
    )


def crop_around_center(image, width, height):
    """
    Given a NumPy / OpenCV 2 image, crops it to the given width and height,
    around it's centre point

    Source: http://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders
    """

    image_size = (image.shape[1], image.shape[0])
    image_center = (int(image_size[0] * 0.5), int(image_size[1] * 0.5))

    if (width > image_size[0]):
        width = image_size[0]

    if (height > image_size[1]):
        height = image_size[1]

    x1 = int(image_center[0] - width * 0.5)
    x2 = int(image_center[0] + width * 0.5)
    y1 = int(image_center[1] - height * 0.5)
    y2 = int(image_center[1] + height * 0.5)

    return image[y1:y2, x1:x2]


def crop_largest_rectangle(image, angle, height, width):
    """
    Crop around the center the largest possible rectangle
    found with largest_rotated_rect.
    """
    return crop_around_center(
        image,
        *largest_rotated_rect(
            width,
            height,
            math.radians(angle)
        )
    )


def generate_rotated_image(image, angle, size=None, crop_center=False,
                           crop_largest_rect=False):
    """
    Generate a valid rotated image for the RotNetDataGenerator. If the
    image is rectangular, the crop_center option should be used to make
    it square. To crop out the black borders after rotation, use the
    crop_largest_rect option. To resize the final image, use the size
    option.
    """
    height, width = image.shape[:2]
    if crop_center:
        if width < height:
            height = width
        else:
            width = height

    image = rotate(image, angle)

    if crop_largest_rect:
        image = crop_largest_rectangle(image, angle, height, width)

    if size:
        image = cv2.resize(image, size)

    return image


def get_filenames(path, num_samples=None):
    image_paths = []

    # Recursively traverse the directory structure
    for root, _, filenames in os.walk(path):
        for filename in filenames:
            image_paths.append(os.path.join(root, filename))

    # Use all available images if num_samples is not provided
    if num_samples is None:
        train_filenames = image_paths
    else:
        # Randomly select num_samples images
        train_filenames = random.sample(
            image_paths, min(num_samples, len(image_paths)))

    # Shuffle the train_filenames list
    random.shuffle(train_filenames)

    return train_filenames

# Dataset class for RotNet data


class RotNetDataset(Dataset):
    def __init__(self, filenames, num_samples, input_shape, transform=None):
        self.filenames = filenames
        self.num_samples = num_samples
        self.input_shape = input_shape
        self.transform = transform

    def __len__(self):
        if self.num_samples is None:
            return len(self.filenames)
        else:
            return min(self.num_samples, len(self.filenames))

    def __getitem__(self, idx):
        image = cv2.imread(self.filenames[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        rotation_angle = random.choice([0, 90, 180, 270])

        # Map rotation angles to class indices (0, 1, 2, 3)
        class_index = rotation_angle // 90  # 0, 1, 2, 3

        rotated_image = generate_rotated_image(
            image,
            rotation_angle,
            size=self.input_shape,
            crop_center=True,
            crop_largest_rect=True
        )

        sample = {'image': rotated_image, 'angle': class_index}

        if self.transform:
            sample['image'] = self.transform(sample['image'])

        return sample


# Define transforms
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize(input_shape),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Run
def setup_training(data_path, validation_path, num_samples, num_val_samples):
    train_filenames = get_filenames(data_path, num_samples)
    validation_filenames = get_filenames(validation_path, num_val_samples)

    print(len(train_filenames), 'train samples')
    print(len(validation_filenames), 'validation samples')

    batch_size = 8
    train_dataset = RotNetDataset(
        train_filenames, num_samples, input_shape, transform=transform)
    train_dataloader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=True, num_workers=1)
    validation_Data = RotNetDataset(
        validation_filenames, num_samples, input_shape, transform=transform)
    validation_dataloader = DataLoader(
        validation_Data, batch_size=batch_size, shuffle=False, num_workers=1)

    model = RotNetModel(num_classes=num_classes).to(device)

    optimizer = optim.SGD(
        model.parameters(),
        lr=0.01,  # Default learning rate
        momentum=0.9,  # Default momentum
    )
    return train_dataloader, validation_dataloader, model, optimizer

def train_model(config, checkpoint_dir=None):
    train_dataloader, validation_dataloader, model, optimizer = setup_training(
        data_path, validation_path, num_samples, num_val_samples)
    criterion = torch.nn.CrossEntropyLoss()

    if checkpoint_dir:
        checkpoint_dict = session.get_checkpoint()
        if checkpoint_dict:
            model.load_state_dict(checkpoint_dict["model_state_dict"])
            optimizer.load_state_dict(checkpoint_dict["optimizer_state_dict"])

    step = 1

    while True:
        model.train()
        running_loss = 0.0

        for i, data in enumerate(train_dataloader, 0):
            inputs, labels = data['image'].to(device), data['angle'].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        epoch_loss = running_loss / len(train_dataloader)

        # Placeholder accuracy calculation
        acc = 0.0

        checkpoint = Checkpoint.from_config(checkpoint_dir)
        if step % config["checkpoint_interval"] == 0:
            checkpoint = Checkpoint.from_dict({
                "step": step,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
            })

        session.report(
            {"mean_accuracy": acc, "lr": config["lr"]},
            checkpoint=checkpoint
        )
        step += 1

perturbation_interval = 5
scheduler = PopulationBasedTraining(
    time_attr="training_iteration",
    perturbation_interval=perturbation_interval,
    metric="mean_accuracy",
    mode="max",
    hyperparam_mutations={
        # distribution for resampling
        "lr": tune.uniform(0.0001, 1),
        # allow perturbations within this set of categorical values
        "momentum": [0.8, 0.9, 0.99],
    },
)

if ray.is_initialized():
    ray.shutdown()
ray.init()

tuner = tune.Tuner(
    train_model,
    run_config=air.RunConfig(
        name="pbt_test",
        # Stop when we've reached a threshold accuracy, or a maximum
        # training_iteration, whichever comes first
        stop={"mean_accuracy": 0.96, "training_iteration": 50},
        verbose=1,
        checkpoint_config=air.CheckpointConfig(
            checkpoint_score_attribute="mean_accuracy",
            num_to_keep=4,
        ),
        storage_path="/tmp/ray_results",
    ),
    tune_config=tune.TuneConfig(
        scheduler=scheduler,
        num_samples=4,
    ),
    param_space={
        "lr": tune.uniform(0.001, 1),
        "momentum": tune.uniform(0.001, 1),
        "checkpoint_interval": perturbation_interval,
    },
)

results_grid = tuner.fit()