import torch.nn.functional as F
import numpy as np
import torch
from torchvision.transforms._functional_tensor import normalize
import pandas as pd
from skimage.segmentation import expand_labels
from skimage.measure import label
import albumentations as A
from monai.inferers import SlidingWindowInferer
from dnafiber.deployment import _get_model
from dnafiber.postprocess import refine_segmentation

transform = A.Compose(
    [
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        A.ToTensorV2(),
    ]
)


def preprocess_image(image):
    image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0)
    image = normalize(image, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    return image


def convert_to_dataset(counts):
    data = {"index": [], "red": [], "green": [], "ratio": []}
    for k, v in counts.items():
        data["index"].append(k)
        data["green"].append(v["green"])
        data["red"].append(v["red"])
        if v["red"] == 0:
            data["ratio"].append(np.nan)
        else:
            data["ratio"].append(v["green"] / (v["red"]))
    df = pd.DataFrame(data)
    return df


def convert_mask_to_image(mask, expand=False):
    if expand:
        mask = expand_labels(mask, distance=expand)
    h, w = mask.shape
    image = np.zeros((h, w, 3), dtype=np.uint8)
    GREEN = np.array([0, 255, 0])
    RED = np.array([255, 0, 0])

    image[mask == 1] = RED
    image[mask == 2] = GREEN

    return image


@torch.inference_mode()
def infer(model, image, device, scale=0.13, to_numpy=True, only_probabilities=False):
    if isinstance(model, str):
        model = _get_model(device=device, revision=model)
    model_pixel_size = 0.26

    scale = scale / model_pixel_size
    tensor = transform(image=image)["image"].unsqueeze(0).to(device)
    h, w = tensor.shape[2], tensor.shape[3]
    device = torch.device(device)
    with torch.autocast(device_type=device.type):
        tensor = F.interpolate(
            tensor,
            size=(int(h * scale), int(w * scale)),
            mode="bilinear",
        )
        if tensor.shape[2] > 1024 or tensor.shape[3] > 1024:
            inferer = SlidingWindowInferer(
                roi_size=(1024, 1024),
                sw_batch_size=4,
                overlap=0.25,
                mode="gaussian",
                device=device,
                progress=True,
            )
            output = inferer(tensor, model)
        else:
            output = model(tensor)

        probabilities = F.softmax(output, dim=1)
        if only_probabilities:
            probabilities = probabilities.cpu()

            probabilities = F.interpolate(
                probabilities,
                size=(h, w),
                mode="bilinear",
            )
            return probabilities

        output = F.interpolate(
            probabilities.argmax(dim=1, keepdim=True).float(),
            size=(h, w),
            mode="nearest",
        )

    output = output.squeeze().byte()
    if to_numpy:
        output = output.cpu().numpy()

    return output