"""Find your apparel effortlessly. Just describe your apparel and get the relevant recommendations with links!"""
import argparse
import logging
import os
import wandb
import gradio as gr

import zipfile
import pickle
from pathlib import Path
from typing import List, Any, Dict
from PIL import Image
from pathlib import Path

from transformers import AutoTokenizer
from sentence_transformers import SentenceTransformer, util
from multilingual_clip import pt_multilingual_clip
import torch

from pathlib import Path
from typing import Callable, Dict, List, Tuple
from PIL.Image import Image

print(__file__)

os.environ["CUDA_VISIBLE_DEVICES"] = ""  # do not use GPU

logging.basicConfig(level=logging.INFO)
DEFAULT_APPLICATION_NAME = "fashion-aggregator"

APP_DIR = Path(__file__).resolve().parent  # what is the directory for this application?
LOGO = APP_DIR / "1001epochs.png"  # path to a small image for display in browser tab and social media
README = APP_DIR / "README.md"  # path to an app readme file in HTML/markdown

DEFAULT_PORT = 11700

EMBEDDINGS_DIR = "artifacts/img-embeddings"
EMBEDDINGS_FILE = os.path.join(EMBEDDINGS_DIR, "embeddings.pkl")
RAW_PHOTOS_DIR = "artifacts/raw-photos"

# Download image embeddings and raw photos
wandb.login(key="4b5a23a662b20fdd61f2aeb5032cf56fdce278a4")  # os.getenv('wandb')
api = wandb.Api()
artifact_embeddings = api.artifact("ryparmar/fashion-aggregator/unimoda-images:v1")
artifact_embeddings.download(EMBEDDINGS_DIR)
artifact_raw_photos = api.artifact("ryparmar/fashion-aggregator/unimoda-raw-images:v1")
artifact_raw_photos.download("artifacts")

with zipfile.ZipFile("artifacts/unimoda.zip", 'r') as zip_ref:
    zip_ref.extractall(RAW_PHOTOS_DIR)


class TextEncoder:
    """Encodes the given text"""

    def __init__(self, model_path="M-CLIP/XLM-Roberta-Large-Vit-B-32"):
        self.model = pt_multilingual_clip.MultilingualCLIP.from_pretrained(model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)

    @torch.no_grad()
    def encode(self, query: str) -> torch.Tensor:
        """Predict/infer text embedding for a given query."""
        query_emb = self.model.forward([query], self.tokenizer)
        return query_emb


class ImageEnoder:
    """Encodes the given image"""

    def __init__(self, model_path="clip-ViT-B-32"):
        self.model = SentenceTransformer(model_path)

    @torch.no_grad()
    def encode(self, image: Image) -> torch.Tensor:
        """Predict/infer text embedding for a given query."""
        image_emb = self.model.encode([image], convert_to_tensor=True, show_progress_bar=False)
        return image_emb


class Retriever:
    """Retrieves relevant images for a given text embedding."""

    def __init__(self, image_embeddings_path=None):
        self.text_encoder = TextEncoder()
        self.image_encoder = ImageEnoder()

        with open(image_embeddings_path, "rb") as file:
            self.image_names, self.image_embeddings = pickle.load(file)
            self.image_names = [
                img_name.replace("fashion-aggregator/fashion_aggregator/data/photos/", "")
                for img_name in self.image_names
            ]
        print("Images:", len(self.image_names))

    @torch.no_grad()
    def predict(self, text_query: str, k: int = 10) -> List[Any]:
        """Return top-k relevant items for a given embedding"""
        query_emb = self.text_encoder.encode(text_query)
        relevant_images = util.semantic_search(query_emb, self.image_embeddings, top_k=k)[0]
        return relevant_images

    @torch.no_grad()
    def search_images(self, text_query: str, k: int = 6) -> Dict[str, List[Any]]:
        """Return top-k relevant images for a given embedding"""
        images = self.predict(text_query, k)
        paths_and_scores = {"path": [], "score": []}
        for img in images:
            paths_and_scores["path"].append(os.path.join(RAW_PHOTOS_DIR, self.image_names[img["corpus_id"]]))
            paths_and_scores["score"].append(img["score"])
        return paths_and_scores


def main(args):
    predictor = PredictorBackend(url=args.model_url)
    frontend = make_frontend(predictor.run, flagging=args.flagging, gantry=args.gantry, app_name=args.application)
    frontend.launch(
        # server_name="0.0.0.0",  # make server accessible, binding all interfaces  # noqa: S104
        # server_port=args.port,  # set a port to bind to, failing if unavailable
        # share=False,  # should we create a (temporary) public link on https://gradio.app?
        # favicon_path=FAVICON,  # what icon should we display in the address bar?
    )


def make_frontend(
    fn: Callable[[Image], str], flagging: bool = False, gantry: bool = False, app_name: str = "fashion-aggregator"
):
    """Creates a gradio.Interface frontend for text to image search function."""

    allow_flagging = "never"

    # Create a Blocks interface
    frontend = gr.Interface(
        fn=fn,
        inputs=gr.inputs.Textbox(label="Item Description", placeholder="Enter item description here"),
        outputs=gr.outputs.Gallery(label="Relevant Items", show_labels=True, label_position="below"),
        title="<span class='title-text'>Fashion Aggregator</span><img src='{logo}' class='logo-image'>".format(logo=LOGO),
        thumbnail=FAVICON,
        description="Discover your perfect apparel effortlessly. Simply describe what you're looking for!",
        cache_examples=False,
        allow_flagging=allow_flagging,
        flagging_options=["incorrect", "offensive", "other"],
    )

    # Customize the HTML template to include the disclaimer and style the title
    frontend.template = """
    <!DOCTYPE html>
    <html>
    <head>
        <style>
            /* Modify the styles according to your preferences */
            .gradio-interface-container {
                text-align: center;
            }
            .gradio-input-container {
                margin-bottom: 20px;
            }
            .gradio-output-container {
                margin-top: 20px;
            }
            .disclaimer {
                font-size: 12px;
                color: #777777;
                margin-top: 20px;
                text-align: center;
            }
            .title-text {
                display: inline-block;
            }
            .logo-image {
                display: inline-block;
                vertical-align: middle;
                margin-left: 10px;
            }
        </style>
    </head>
    <body>
        {{interface_html}}
        <div class="disclaimer">
            <gr-blocks>
                <gr-row>
                    <gr-column>
                        <span>Disclaimer: This web app is for demonstration purposes only and not intended for commercial use.</span>
                    </gr-column>
                </gr-row>
            </gr-blocks>
        </div>
    </body>
    </html>
    """

    return frontend

# def make_frontend(
#     fn: Callable[[Image], str], flagging: bool = False, gantry: bool = False, app_name: str = "fashion-aggregator"
# ):
#     """Creates a gradio.Interface frontend for text to image search function."""

#     allow_flagging = "never"

#     # Build a customized browser interface to a Python function
#     frontend = gr.Interface(
#         fn=fn,
#         outputs=gr.Gallery(label="Relevant Items", show_labels=True, label_position="below"),
#         inputs=gr.inputs.Textbox(label="Item Description", placeholder="Enter item description here"),
#         title="Fashion Aggregator",  
#         thumbnail=FAVICON,
#         description="Discover your perfect apparel effortlessly. Simply describe what you're looking for!",
#         cache_examples=False,
#         allow_flagging=allow_flagging,
#         flagging_options=["incorrect", "offensive", "other"],
#     )

#     return frontend


class PredictorBackend:
    """Interface to a backend that serves predictions.

    To communicate with a backend accessible via a URL, provide the url kwarg.

    Otherwise, runs a predictor locally.
    """

    def __init__(self, url=None):
        if url is not None:
            self.url = url
            self._predict = self._predict_from_endpoint
        else:
            model = Retriever(image_embeddings_path=EMBEDDINGS_FILE)
            self._predict = model.predict
            self._search_images = model.search_images

    def run(self, text: str):
        pred, metrics = self._predict_with_metrics(text)
        self._log_inference(pred, metrics)
        return pred

    def _predict_with_metrics(self, text: str) -> Tuple[List[str], Dict[str, float]]:
        paths_and_scores = self._search_images(text)
        metrics = {"mean_score": sum(paths_and_scores["score"]) / len(paths_and_scores["score"])}
        return paths_and_scores["path"], metrics

    def _log_inference(self, pred, metrics):
        for key, value in metrics.items():
            logging.info(f"METRIC {key} {value}")
        logging.info(f"PRED >begin\n{pred}\nPRED >end")


def _make_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model_url",
        default=None,
        type=str,
        help="Identifies a URL to which to send image data. Data is base64-encoded, converted to a utf-8 string, and then set via a POST request as JSON with the key 'image'. Default is None, which instead sends the data to a model running locally.",
    )
    parser.add_argument(
        "--port",
        default=DEFAULT_PORT,
        type=int,
        help=f"Port on which to expose this server. Default is {DEFAULT_PORT}.",
    )
    parser.add_argument(
        "--flagging",
        action="store_true",
        help="Pass this flag to allow users to 'flag' model behavior and provide feedback.",
    )
    parser.add_argument(
        "--gantry",
        action="store_true",
        help="Pass --flagging and this flag to log user feedback to Gantry. Requires GANTRY_API_KEY to be defined as an environment variable.",
    )
    parser.add_argument(
        "--application",
        default=DEFAULT_APPLICATION_NAME,
        type=str,
        help=f"Name of the Gantry application to which feedback should be logged, if --gantry and --flagging are passed. Default is {DEFAULT_APPLICATION_NAME}.",
    )
    return parser


if __name__ == "__main__":
    parser = _make_parser()
    args = parser.parse_args()
    main(args)