import pickle
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import tifffile as tiff
import torch
import torch.nn.functional as F

from model import Discriminator, Generator


class StarNet:
    def __init__(
        self,
        mode,
        window_size=512,
        stride=256,
        lr=1e-4,
        train_folder="./train/",
        batch_size=1,
        device=None,
    ):
        if mode not in ("RGB", "Greyscale"):
            raise ValueError("Mode should be either RGB or Greyscale")
        if window_size % 256 != 0:
            raise ValueError("window_size must be divisible by 256")
        if stride <= 0 or stride > window_size:
            raise ValueError("stride must be greater than zero and no larger than window_size")
        if (window_size - stride) % 2 != 0:
            raise ValueError("window_size - stride must be even")

        self.mode = mode
        self.input_channels = 3 if mode == "RGB" else 1
        self.window_size = window_size
        self.stride = stride
        self.lr = lr
        self.train_folder = Path(train_folder)
        self.batch_size = batch_size
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

        self.history = {}
        self._ema = 0.9999
        self.weights = []
        self.original = []
        self.starless = []
        self.iters_per_epoch = 0

        self.G = None
        self.D = None
        self.gen_optimizer = None
        self.dis_optimizer = None

    def __str__(self):
        return "StarNet PyTorch instance"

    def load_training_dataset(self):
        original_dir = self.train_folder / "original"
        starless_dir = self.train_folder / "starless"
        original_files = sorted(path.name for path in original_dir.glob("*.tif"))
        starless_files = sorted(path.name for path in starless_dir.glob("*.tif"))

        if not original_files or not starless_files:
            raise ValueError("No training data found in %s" % self.train_folder)
        if original_files != starless_files:
            raise ValueError(
                "Corresponding names in the original and starless folders must be equal"
            )

        print("Total training images found: %d" % len(original_files))
        self.original = [tiff.imread(original_dir / name) for name in original_files]
        self.starless = [tiff.imread(starless_dir / name) for name in starless_files]
        self.weights = []
        total_pixels = 0

        for name, original, starless in zip(original_files, self.original, self.starless):
            if original.shape != starless.shape:
                raise ValueError("Image sizes are not equal for %s" % name)
            if original.shape[0] < self.window_size or original.shape[1] < self.window_size:
                raise ValueError(
                    "%s is smaller than the configured window size %d"
                    % (name, self.window_size)
                )
            pixels = original.shape[0] * original.shape[1]
            total_pixels += pixels
            self.weights.append(pixels)

        self.iters_per_epoch = max(1, total_pixels // (self.window_size * self.window_size))
        self.weights = np.asarray(self.weights, dtype=np.float64)
        self.weights /= self.weights.sum()

        print("Total size of training images: %.2f MP" % (total_pixels / 1e6))
        print("One epoch is set to %d iterations" % self.iters_per_epoch)
        print("Training dataset has been successfully loaded!")

    def _generator(self, m=64):
        if m != 64:
            raise ValueError("The translated architecture supports m=64")
        return Generator(input_channels=self.input_channels, batch_norm_eps=1e-3)

    def _discriminator(self):
        return Discriminator(input_channels=self.input_channels, batch_norm_eps=1e-3)

    def _weight_paths(self, weights):
        path = Path(weights)
        
        if path.suffix.lower() in (".pth", ".pt"):
            return path, None
        return (
            Path(str(path) + "_G_" + self.mode + ".pth"),
            Path(str(path) + "_D_" + self.mode + ".pth"),
        )

    @staticmethod
    def _load_state(module, path, device):
        try:
            checkpoint = torch.load(path, map_location=device, weights_only=True)
        except TypeError:  # PyTorch < 2.0 compatibility
            checkpoint = torch.load(path, map_location=device)
        state_dict = checkpoint.get("state_dict", checkpoint)
        module.load_state_dict(state_dict)

    def load_model(self, weights=None, history=None):
        self.G = self._generator().to(self.device)
        self.D = self._discriminator().to(self.device)
        self.gen_optimizer = torch.optim.Adam(self.G.parameters(), lr=self.lr)
        self.dis_optimizer = torch.optim.Adam(self.D.parameters(), lr=self.lr / 4)

        if weights:
            generator_path, discriminator_path = self._weight_paths(weights)
            if not generator_path.exists():
                raise FileNotFoundError("Generator weights were not found at %s" % generator_path)
            self._load_state(self.G, generator_path, self.device)

            if discriminator_path and discriminator_path.exists():
                self._load_state(self.D, discriminator_path, self.device)
            elif discriminator_path:
                print("Discriminator weights were not found; using initialized weights.")

        if history:
            history_path = Path(str(history) + "_" + self.mode + ".pkl")
            with history_path.open("rb") as handle:
                self.history = pickle.load(handle)

        self.G.eval()
        self.D.eval()
        return self.G

    def initialize_model(self):
        return self.load_model()

    @staticmethod
    def _ramp(x):
        if torch.is_tensor(x):
            return torch.clamp(x, 0.0, 1.0)
        return np.clip(x, 0.0, 1.0)

    @staticmethod
    def _image_scale(image):
        if image.dtype == np.uint8:
            return 255.0
        if image.dtype == np.uint16:
            return 65535.0
        if np.issubdtype(image.dtype, np.floating):
            return 1.0
        raise ValueError("Unsupported training image dtype: %s" % image.dtype)

    def _get_sample(self, r, h, w, type):
        if type not in ("original", "starless"):
            raise ValueError("type must be original or starless")
        source = self.original if type == "original" else self.starless
        sample = source[r][h:h + self.window_size, w:w + self.window_size]
        return sample.astype(np.float32) / self._image_scale(sample)

    def _augmentator(self, original, starless):
        original = np.array(original, copy=True)
        starless = np.array(starless, copy=True)

        if np.random.rand() < 0.5:
            original = np.flip(original, axis=1)
            starless = np.flip(starless, axis=1)
        if np.random.rand() < 0.5:
            original = np.flip(original, axis=0)
            starless = np.flip(starless, axis=0)
        if np.random.rand() < 0.5:
            rotations = np.random.randint(1, 4)
            original = np.rot90(original, rotations, axes=(1, 0))
            starless = np.rot90(starless, rotations, axes=(1, 0))

        if self.mode == "RGB":
            if original.ndim != 3 or original.shape[2] < 3:
                raise ValueError("RGB training images must contain at least three channels")
            original = original[:, :, :3]
            starless = starless[:, :, :3]
            if np.random.rand() < 0.7:
                channel = np.random.randint(3)
                minimum = min(float(original.min()), float(starless.min()))
                offset = np.random.rand() * 0.25 - np.random.rand() * minimum
                original[:, :, channel] += offset * (1.0 - original[:, :, channel])
                starless[:, :, channel] += offset * (1.0 - starless[:, :, channel])
            if np.random.rand() < 0.7:
                sequence = np.random.permutation(3)
                original = original[:, :, sequence]
                starless = starless[:, :, sequence]
        else:
            if np.random.rand() < 0.7:
                minimum = min(float(original.min()), float(starless.min()))
                offset = np.random.rand() * 0.25 - np.random.rand() * minimum
                original += offset * (1.0 - original)
                starless += offset * (1.0 - starless)
            if original.ndim == 3:
                channel = np.random.randint(min(3, original.shape[2]))
                original = original[:, :, channel]
                starless = starless[:, :, channel]
            original = original[:, :, None]
            starless = starless[:, :, None]

        return (
            np.ascontiguousarray(np.clip(original, 0.0, 1.0)),
            np.ascontiguousarray(np.clip(starless, 0.0, 1.0)),
        )

    def generate_input(self, iterations=1, augmentation=False):
        if not self.original:
            raise RuntimeError("Training dataset was not loaded")

        original_batch = None
        starless_batch = None
        for _ in range(iterations):
            original_batch = np.zeros(
                (
                    self.batch_size,
                    self.window_size,
                    self.window_size,
                    self.input_channels,
                ),
                dtype=np.float32,
            )
            starless_batch = np.zeros_like(original_batch)

            for index in range(self.batch_size):
                image_index = int(
                    np.random.choice(len(self.original), p=np.asarray(self.weights))
                )
                image_h, image_w = self.original[image_index].shape[:2]
                top = np.random.randint(0, image_h - self.window_size + 1)
                left = np.random.randint(0, image_w - self.window_size + 1)
                original = self._get_sample(image_index, top, left, "original")
                starless = self._get_sample(image_index, top, left, "starless")

                if augmentation:
                    original, starless = self._augmentator(original, starless)
                elif self.mode == "Greyscale":
                    if original.ndim == 3:
                        channel = np.random.randint(min(3, original.shape[2]))
                        original = original[:, :, channel]
                        starless = starless[:, :, channel]
                    original = original[:, :, None]
                    starless = starless[:, :, None]
                else:
                    original = original[:, :, :3]
                    starless = starless[:, :, :3]

                original_batch[index] = original
                starless_batch[index] = starless

        return original_batch, starless_batch

    @staticmethod
    def _set_requires_grad(module, enabled):
        for parameter in module.parameters():
            parameter.requires_grad_(enabled)

    def _losses(self, original, target):
        generated = self.G(original)
        real_outputs = self.D(target)
        fake_outputs = self.D(generated)

        real_features = real_outputs[:-1]
        fake_features = fake_outputs[:-1]
        predict_real = real_outputs[-1]
        predict_fake = fake_outputs[-1]

        dis_loss = -(
            torch.log(predict_real + 1e-8) + torch.log(1.0 - predict_fake + 1e-8)
        ).mean()
        gen_loss_gan = -torch.log(predict_fake + 1e-8).mean()
        perceptual = [
            torch.mean(torch.abs(fake - real))
            for fake, real in zip(fake_features, real_features)
        ]
        gen_l1 = F.l1_loss(generated, target)
        gen_loss = (
            gen_loss_gan * 0.1
            + perceptual[0] * 0.1
            + sum(loss * 10.0 for loss in perceptual[1:])
            + gen_l1 * 100.0
        )

        metrics = {
            "dis_loss": dis_loss,
            "gen_loss_GAN": gen_loss_gan,
            "gen_p1": perceptual[0],
            "gen_p2": perceptual[1],
            "gen_p3": perceptual[2],
            "gen_p4": perceptual[3],
            "gen_p5": perceptual[4],
            "gen_p6": perceptual[5],
            "gen_p7": perceptual[6],
            "gen_p8": perceptual[7],
            "gen_L1": gen_l1 * 100.0,
            "total": gen_loss,
        }
        return gen_loss, dis_loss, metrics

    def _train_step(self, original, target):
        self.G.train()
        self.D.train()
        self.gen_optimizer.zero_grad(set_to_none=True)
        self.dis_optimizer.zero_grad(set_to_none=True)

        gen_loss, dis_loss, metrics = self._losses(original, target)
        generator_gradients = torch.autograd.grad(
            gen_loss,
            tuple(self.G.parameters()),
            retain_graph=True,
        )
        discriminator_gradients = torch.autograd.grad(
            dis_loss,
            tuple(self.D.parameters()),
        )

        for parameter, gradient in zip(self.G.parameters(), generator_gradients):
            parameter.grad = gradient
        for parameter, gradient in zip(self.D.parameters(), discriminator_gradients):
            parameter.grad = gradient

        self.gen_optimizer.step()
        self.dis_optimizer.step()
        return {name: float(value.detach().cpu()) for name, value in metrics.items()}

    def _update_history(self, metrics):
        for name, value in metrics.items():
            if name in self.history and self.history[name]:
                value = value * (1.0 - self._ema) + self.history[name][-1] * self._ema
                self.history[name].append(value)
            else:
                self.history[name] = [value]

    def _plot_progress(self, original, target):
        self.G.eval()
        with torch.inference_mode():
            generated = self.G(original)

        original = ((original[0] + 1.0) / 2.0).detach().cpu().numpy()
        generated = ((generated[0] + 1.0) / 2.0).detach().cpu().numpy()
        target = ((target[0] + 1.0) / 2.0).detach().cpu().numpy()

        plt.close()
        _, axes = plt.subplots(1, 3, sharex=True, figsize=(16.5, 5.5))
        titles = ("Original", "Starless", "Target")
        for axis, title, image in zip(axes, titles, (original, generated, target)):
            if self.mode == "RGB":
                axis.imshow(np.transpose(image, (1, 2, 0)))
            else:
                axis.imshow(image[0], cmap="gray", vmin=0, vmax=1)
            axis.set_title(title)
        plt.pause(0.001)
        self.G.train()

    def train(
        self,
        epochs,
        augmentation=True,
        plot_progress=False,
        plot_interval=50,
        save_backups=True,
        warm_up=False,
    ):
        if not self.original:
            raise RuntimeError("Training dataset was not loaded; call load_training_dataset()")
        if self.G is None or self.D is None:
            self.initialize_model()

        for epoch in range(epochs):
            for iteration in range(self.iters_per_epoch):
                original, target = self.generate_input(augmentation=augmentation)
                original = torch.from_numpy(original.transpose(0, 3, 1, 2)).to(self.device)
                target = torch.from_numpy(target.transpose(0, 3, 1, 2)).to(self.device)
                original = original * 2.0 - 1.0
                target = target * 2.0 - 1.0
                if warm_up:
                    target = original

                if plot_progress and iteration % plot_interval == 0:
                    self._plot_progress(original, target)

                metrics = self._train_step(original, target)
                self._update_history(metrics)
                print(
                    "\rEpoch: %d. Iteration %d / %d Loss %f    "
                    % (epoch, iteration, self.iters_per_epoch, self.history["total"][-1]),
                    end="",
                )

            print()
            if save_backups:
                suffix = "even" if epoch % 2 == 0 else "odd"
                self._save_network(self.G, "./starnet_backup_G_%s.pth" % suffix)
                self._save_network(self.D, "./starnet_backup_D_%s.pth" % suffix)

        if plot_progress:
            plt.close()
        self.G.eval()
        self.D.eval()
        return self.history

    def plot_history(self, last=None):
        if not self.history:
            raise RuntimeError("Empty training history, nothing to plot")

        keys = list(self.history.keys())
        columns = 3
        rows = int(np.ceil(len(keys) / columns))
        figure, axes = plt.subplots(rows, columns, sharex=True, figsize=(16, 3.5 * rows))
        axes = np.atleast_1d(axes).reshape(-1)
        for axis, key in zip(axes, keys):
            values = self.history[key][-last:] if last else self.history[key]
            axis.plot(values)
            axis.set_title(key)
        for axis in axes[len(keys):]:
            axis.set_visible(False)
        figure.tight_layout()
        return figure, axes

    @staticmethod
    def _save_network(network, path):
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"state_dict": network.state_dict()}, path)

    def save_model(self, weights_filename, history_filename=None):
        if self.G is None or self.D is None:
            raise RuntimeError("Model has not been initialized")

        generator_path = Path(str(weights_filename) + "_G_" + self.mode + ".pth")
        discriminator_path = Path(str(weights_filename) + "_D_" + self.mode + ".pth")
        self._save_network(self.G, generator_path)
        self._save_network(self.D, discriminator_path)

        if history_filename:
            history_path = Path(str(history_filename) + "_" + self.mode + ".pkl")
            history_path.parent.mkdir(parents=True, exist_ok=True)
            with history_path.open("wb") as handle:
                pickle.dump(self.history, handle)
        return generator_path, discriminator_path

    def _read_image(self, input_path):
        data = tiff.imread(input_path)
        if data.ndim > 3:
            layer = input("Tiff has %d layers, please enter layer to process: " % data.shape[0])
            data = data[int(layer)]

        if data.dtype == np.dtype("uint16"):
            scale = 65535.0
        elif data.dtype == np.dtype("uint8"):
            scale = 255.0
        else:
            raise ValueError("Unknown image dtype: %s" % data.dtype)

        image = (data / scale).astype(np.float32)
        if self.mode == "Greyscale":
            if image.ndim != 2:
                raise ValueError("You loaded Greyscale model, but the image is RGB")
            image = image[:, :, None]
        else:
            if image.ndim != 3:
                raise ValueError("You loaded RGB model, but the image is Greyscale")
            if image.shape[2] == 4:
                print("Input image has 4 channels. Removing Alpha-Channel")
                image = image[:, :, :3]
            if image.shape[2] != 3:
                raise ValueError("RGB input must have 3 or 4 channels")
        return image, data.dtype, scale

    @staticmethod
    def _imwrite(path, image):
        if hasattr(tiff, "imwrite"):
            tiff.imwrite(path, image)
        else:
            tiff.imsave(path, image)

    def transform(self, in_name, out_name):
        if self.G is None:
            raise RuntimeError("Call load_model() before transform()")

        self.G.eval()
        image, input_dtype, scale = self._read_image(in_name)
        offset = (self.window_size - self.stride) // 2
        h, w, _ = image.shape
        ith = h // self.stride + 1
        itw = w // self.stride + 1
        dh = ith * self.stride - h
        dw = itw * self.stride - w

        image = np.concatenate((image, image[(h - dh):, :, :]), axis=0)
        image = np.concatenate((image, image[:, (w - dw):, :]), axis=1)
        padded_h, padded_w, _ = image.shape
        image = np.concatenate((image, image[(padded_h - offset):, :, :]), axis=0)
        image = np.concatenate((image[:offset, :, :], image), axis=0)
        image = np.concatenate((image, image[:, (padded_w - offset):, :]), axis=1)
        image = np.concatenate((image[:, :offset, :], image), axis=1)
        image = image * 2.0 - 1.0
        output = np.empty_like(image)

        with torch.inference_mode():
            for i in range(ith):
                for j in range(itw):
                    x = self.stride * i
                    y = self.stride * j
                    tile = image[
                        x:x + self.window_size,
                        y:y + self.window_size,
                        :,
                    ]
                    tile = (
                        torch.from_numpy(tile.transpose(2, 0, 1))
                        .unsqueeze(0)
                        .to(self.device)
                    )
                    result = self.G(tile)
                    result = ((result[0] + 1.0) / 2.0).cpu().numpy().transpose(1, 2, 0)
                    output[
                        x + offset:x + offset + self.stride,
                        y + offset:y + offset + self.stride,
                        :,
                    ] = result[
                        offset:offset + self.stride,
                        offset:offset + self.stride,
                        :,
                    ]
                    progress = int((itw * i + j + 1) * 100 / (ith * itw))
                    print("Transforming input image... %d%%\r" % progress, end="")

        print("Transforming input image... Done!")
        output = np.clip(output, 0.0, 1.0)
        output = output[offset:offset + h, offset:offset + w, :]
        if self.mode == "Greyscale":
            output = output[:, :, 0]

        output_path = Path(out_name)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        self._imwrite(output_path, (output * scale).astype(input_dtype))
        return output_path