import os
import sys
from collections import namedtuple
from pathlib import Path
import re
import torch
import torch.hub # pylint: disable=ungrouped-imports
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from modules import devices, paths, shared, lowvram, errors


blip_image_eval_size = 384
clip_model_name = 'ViT-L/14'
Category = namedtuple("Category", ["name", "topn", "items"])
re_topn = re.compile(r"\.top(\d+)\.")


def category_types():
    return [f.stem for f in Path(shared.interrogator.content_dir).glob('*.txt')]


def download_default_clip_interrogate_categories(content_dir):
    shared.log.info("Downloading CLIP categories...")
    tmpdir = f"{content_dir}_tmp"
    cat_types = ["artists", "flavors", "mediums", "movements"]
    try:
        os.makedirs(tmpdir, exist_ok=True)
        for category_type in cat_types:
            torch.hub.download_url_to_file(f"https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/clip_interrogator/data/{category_type}.txt", os.path.join(tmpdir, f"{category_type}.txt"))
        os.rename(tmpdir, content_dir)
    except Exception as e:
        errors.display(e, "downloading default CLIP interrogate categories")
    finally:
        if os.path.exists(tmpdir):
            os.removedirs(tmpdir)


class InterrogateModels:
    blip_model = None
    clip_model = None
    clip_preprocess = None
    dtype = None
    running_on_cpu = None

    def __init__(self, content_dir):
        self.loaded_categories = None
        self.skip_categories = []
        self.content_dir = content_dir
        self.running_on_cpu = devices.device_interrogate == torch.device("cpu")

    def categories(self):
        if not os.path.exists(self.content_dir):
            download_default_clip_interrogate_categories(self.content_dir)
        if self.loaded_categories is not None and self.skip_categories == shared.opts.interrogate_clip_skip_categories:
            return self.loaded_categories
        self.loaded_categories = []

        if os.path.exists(self.content_dir):
            self.skip_categories = shared.opts.interrogate_clip_skip_categories
            cat_types = []
            for filename in Path(self.content_dir).glob('*.txt'):
                cat_types.append(filename.stem)
                if filename.stem in self.skip_categories:
                    continue
                m = re_topn.search(filename.stem)
                topn = 1 if m is None else int(m.group(1))
                with open(filename, "r", encoding="utf8") as file:
                    lines = [x.strip() for x in file.readlines()]
                self.loaded_categories.append(Category(name=filename.stem, topn=topn, items=lines))
        return self.loaded_categories

    def create_fake_fairscale(self):
        class FakeFairscale:
            def checkpoint_wrapper(self):
                pass
        sys.modules["fairscale.nn.checkpoint.checkpoint_activations"] = FakeFairscale

    def load_blip_model(self):
        self.create_fake_fairscale()
        from repositories.blip import models # pylint: disable=unused-import
        from repositories.blip.models import blip
        import modules.modelloader as modelloader
        model_path = os.path.join(paths.models_path, "BLIP")
        download_name='model_base_caption_capfilt_large.pth',
        shared.log.debug(f'Model interrogate load: type=BLiP model={download_name} path={model_path}')
        files = modelloader.load_models(
            model_path=model_path,
            model_url='https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth',
            ext_filter=[".pth"],
            download_name=download_name,
        )
        blip_model = blip.blip_decoder(pretrained=files[0], image_size=blip_image_eval_size, vit='base', med_config=os.path.join(paths.paths["BLIP"], "configs", "med_config.json")) # pylint: disable=c-extension-no-member
        blip_model.eval()

        return blip_model

    def load_clip_model(self):
        shared.log.debug(f'Model interrogate load: type=CLiP model={clip_model_name} path={shared.opts.clip_models_path}')
        import clip
        if self.running_on_cpu:
            model, preprocess = clip.load(clip_model_name, device="cpu", download_root=shared.opts.clip_models_path)
        else:
            model, preprocess = clip.load(clip_model_name, download_root=shared.opts.clip_models_path)
        model.eval()
        model = model.to(devices.device_interrogate)
        return model, preprocess

    def load(self):
        if self.blip_model is None:
            self.blip_model = self.load_blip_model()
            if not shared.opts.no_half and not self.running_on_cpu:
                self.blip_model = self.blip_model.half()
        self.blip_model = self.blip_model.to(devices.device_interrogate)
        if self.clip_model is None:
            self.clip_model, self.clip_preprocess = self.load_clip_model()
            if not shared.opts.no_half and not self.running_on_cpu:
                self.clip_model = self.clip_model.half()
        self.clip_model = self.clip_model.to(devices.device_interrogate)
        self.dtype = next(self.clip_model.parameters()).dtype

    def send_clip_to_ram(self):
        if not shared.opts.interrogate_keep_models_in_memory:
            if self.clip_model is not None:
                self.clip_model = self.clip_model.to(devices.cpu)

    def send_blip_to_ram(self):
        if not shared.opts.interrogate_keep_models_in_memory:
            if self.blip_model is not None:
                self.blip_model = self.blip_model.to(devices.cpu)

    def unload(self):
        self.send_clip_to_ram()
        self.send_blip_to_ram()
        devices.torch_gc()

    def rank(self, image_features, text_array, top_count=1):
        import clip
        devices.torch_gc()
        if shared.opts.interrogate_clip_dict_limit != 0:
            text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]
        top_count = min(top_count, len(text_array))
        text_tokens = clip.tokenize(list(text_array), truncate=True).to(devices.device_interrogate)
        text_features = self.clip_model.encode_text(text_tokens).type(self.dtype)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        similarity = torch.zeros((1, len(text_array))).to(devices.device_interrogate)
        for i in range(image_features.shape[0]):
            similarity += (100.0 * image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)
        similarity /= image_features.shape[0]
        top_probs, top_labels = similarity.cpu().topk(top_count, dim=-1)
        return [(text_array[top_labels[0][i].numpy()], (top_probs[0][i].numpy()*100)) for i in range(top_count)]

    def generate_caption(self, pil_image):
        gpu_image = transforms.Compose([
            transforms.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=InterpolationMode.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        ])(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)
        with devices.inference_context():
            caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=shared.opts.interrogate_clip_min_length, max_length=shared.opts.interrogate_clip_max_length)
        return caption[0]

    def interrogate(self, pil_image):
        res = ""
        shared.state.begin('interrogate')
        try:
            if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
                lowvram.send_everything_to_cpu()
                devices.torch_gc()
            self.load()
            if isinstance(pil_image, list):
                pil_image = pil_image[0]
            if isinstance(pil_image, dict) and 'name' in pil_image:
                pil_image = Image.open(pil_image['name'])
            pil_image = pil_image.convert("RGB")
            caption = self.generate_caption(pil_image)
            self.send_blip_to_ram()
            devices.torch_gc()
            res = caption
            clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device_interrogate)
            with devices.inference_context(), devices.autocast():
                image_features = self.clip_model.encode_image(clip_image).type(self.dtype)
                image_features /= image_features.norm(dim=-1, keepdim=True)
                for _name, topn, items in self.categories():
                    matches = self.rank(image_features, items, top_count=topn)
                    for match, score in matches:
                        if shared.opts.interrogate_return_ranks:
                            res += f", ({match}:{score/100:.3f})"
                        else:
                            res += f", {match}"
        except Exception as e:
            errors.display(e, 'interrogate')
            res += "<error>"
        self.unload()
        shared.state.end()
        return res

# --------- interrrogate ui

ci = None
low_vram = False


class BatchWriter:
    def __init__(self, folder):
        self.folder = folder
        self.csv, self.file = None, None

    def add(self, file, prompt):
        txt_file = os.path.splitext(file)[0] + ".txt"
        with open(os.path.join(self.folder, txt_file), 'w', encoding='utf-8') as f:
            f.write(prompt)

    def close(self):
        if self.file is not None:
            self.file.close()


def get_clip_models():
    import open_clip
    return ['/'.join(x) for x in open_clip.list_pretrained()]


def load_interrogator(model):
    from clip_interrogator import Config, Interrogator
    global ci # pylint: disable=global-statement
    if ci is None:
        config = Config(device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'Interrogator'), clip_model_name=model, quiet=True)
        if low_vram:
            config.apply_low_vram_defaults()
        shared.log.info(f'Interrogate load: config={config}')
        ci = Interrogator(config)
    elif model != ci.config.clip_model_name:
        ci.config.clip_model_name = model
        shared.log.info(f'Interrogate load: config={ci.config}')
        ci.load_clip_model()


def unload_clip_model():
    if ci is not None:
        shared.log.debug('Interrogate offload')
        ci.caption_model = ci.caption_model.to(devices.cpu)
        ci.clip_model = ci.clip_model.to(devices.cpu)
        ci.caption_offloaded = True
        ci.clip_offloaded = True
        devices.torch_gc()


def interrogate(image, mode, caption=None):
    shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
    if mode == 'best':
        prompt = ci.interrogate(image, caption=caption)
    elif mode == 'caption':
        prompt = ci.generate_caption(image) if caption is None else caption
    elif mode == 'classic':
        prompt = ci.interrogate_classic(image, caption=caption)
    elif mode == 'fast':
        prompt = ci.interrogate_fast(image, caption=caption)
    elif mode == 'negative':
        prompt = ci.interrogate_negative(image)
    else:
        raise RuntimeError(f"Unknown mode {mode}")
    return prompt


def interrogate_image(image, model, mode):
    shared.state.begin()
    shared.state.job = 'interrogate'
    try:
        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
            lowvram.send_everything_to_cpu()
            devices.torch_gc()
        load_interrogator(model)
        image = image.convert('RGB')
        shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
        prompt = interrogate(image, mode)
    except Exception as e:
        prompt = f"Exception {type(e)}"
        shared.log.error(f'Interrogate: {e}')
    shared.state.end()
    return prompt


def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
    files = []
    if batch_files is not None:
        files += [f.name for f in batch_files]
    if batch_folder is not None:
        files += [f.name for f in batch_folder]
    if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
        files += [os.path.join(batch_str, f) for f in os.listdir(batch_str) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    if len(files) == 0:
        shared.log.error('Interrogate batch no images')
        return ''
    shared.state.begin()
    shared.state.job = 'batch interrogate'
    prompts = []
    try:
        if shared.backend == shared.Backend.ORIGINAL and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
            lowvram.send_everything_to_cpu()
            devices.torch_gc()
        load_interrogator(model)
        shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
        captions = []
        # first pass: generate captions
        for file in files:
            caption = ""
            try:
                if shared.state.interrupted:
                    break
                image = Image.open(file).convert('RGB')
                caption = ci.generate_caption(image)
            except Exception as e:
                shared.log.error(f'Interrogate caption: {e}')
            finally:
                captions.append(caption)
        # second pass: interrogate
        if write:
            writer = BatchWriter(os.path.dirname(files[0]))
        for idx, file in enumerate(files):
            try:
                if shared.state.interrupted:
                    break
                image = Image.open(file).convert('RGB')
                prompt = interrogate(image, mode, caption=captions[idx])
                prompts.append(prompt)
                if write:
                    writer.add(file, prompt)
            except OSError as e:
                shared.log.error(f'Interrogate batch: {e}')
        if write:
            writer.close()
        ci.config.quiet = False
        unload_clip_model()
    except Exception as e:
        shared.log.error(f'Interrogate batch: {e}')
    shared.state.end()
    return '\n\n'.join(prompts)


def analyze_image(image, model):
    load_interrogator(model)
    image = image.convert('RGB')
    image_features = ci.image_to_features(image)
    top_mediums = ci.mediums.rank(image_features, 5)
    top_artists = ci.artists.rank(image_features, 5)
    top_movements = ci.movements.rank(image_features, 5)
    top_trendings = ci.trendings.rank(image_features, 5)
    top_flavors = ci.flavors.rank(image_features, 5)
    medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
    artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
    movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
    trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
    flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
    return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks