import torch
import torch.nn.functional as F
import numpy as np

MOLD_HIGH = 0.85
MOLD_LOW  = 0.50
BIO_TH    = 0.70
UNCERT_TH = 0.15
PATCH_TH  = 0.50
DINO_TH   = 0.75

def enable_dropout(model):
    for m in model.modules():
        if m.__class__.__name__.startswith("Dropout"):
            m.train()

def mc_uncertainty(model, img_tensor, mold_idx, T=15):
    enable_dropout(model)
    probs = []
    for _ in range(T):
        out = model(img_tensor.unsqueeze(0))
        cp = F.softmax(out["class"],1)[0]
        probs.append(cp[mold_idx].item())
    model.eval()
    return float(np.mean(probs)), float(np.std(probs))

def patch_consistency(model, image, transform, mold_idx, device, stride=112):
    w,h = image.size
    votes = []
    for y in range(0, h-224+1, stride):
        for x in range(0, w-224+1, stride):
            patch = image.crop((x,y,x+224,y+224))
            pt = transform(patch).unsqueeze(0).to(device)
            with torch.no_grad():
                out = model(pt)
                cp = F.softmax(out["class"],1)[0]
            votes.append(cp[mold_idx].item())
    return float(np.mean(np.array(votes) > 0.7)) if votes else 0.0

def final_decision_v2(
    mold_p, bio_p, uncert, patch_ratio, dino_sim
):
    if (
        mold_p > MOLD_HIGH and
        bio_p > BIO_TH and
        uncert < UNCERT_TH and
        patch_ratio > PATCH_TH and
        dino_sim > DINO_TH
    ):
        return "Mold"

    if mold_p > MOLD_LOW and bio_p > BIO_TH:
        return "Possible Mold"

    return "Not Mold"