from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from pyvis.network import Network
from pprint import pprint
import networkx as nx
import gradio as gr
import re
import datasets
from huggingface_hub import login, HfApi
from datasets import Dataset, load_dataset
from rapidfuzz import fuzz, process
import math
import pandas as pd
import gspread
import torch
import json
from typing import Callable, Optional
from dataclasses import dataclass
from datasets import load_dataset
from transformers import (
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline
)
from peft import PeftModel, LoraConfig, get_peft_model, TaskType

# Setup

REPO_ID_NEAR_FIELD_RAW = "milistu/AMAZON-Products-2023"
REPO_ID_NEAR_FIELD = "aslan-ng/amazon_products_2023"
REPO_ID_FAR_FIELD = "aslan-ng/amazon_products_2025"
REPO_ID_LORA_GREEN_PATENTS = "aslan-ng/lora-green-patents"

def product_quality_score(average_rating: float, rating_number: int):
    """
    Bayesian Average (Amazon-style)
    Args:
      avg_rating: product's average rating
      rating_number: number of reviews
    """
    m = 1  # Minimum number of reviews required (tunable)
    C = 3.5  # Global average rating (baseline)
    if rating_number <= 0 or average_rating is None:
        return C  # fallback to global mean
    return (rating_number / (rating_number + m)) * average_rating + (m / (rating_number + m)) * C

def load_near_field_raw_from_huggingface():
    """
    Load the raw near-field dataset from HuggingFace.
    """
    ds = datasets.load_dataset(REPO_ID_NEAR_FIELD_RAW, split="train")
    print("Initial size: ", len(ds))

    # Drop the extra categories
    main_categories_to_remove = ["meta_Books", "meta_CDs_and_Vinyl", "meta_Digital_Music", "meta_Gift_Cards", "meta_Grocery_and_Gourmet_Food",
                                 "meta_Magazine_Subscriptions", "meta_Software", "meta_Video_Games"]
    ds = ds.filter(lambda row: row["filename"] not in main_categories_to_remove) ###

    # Keep only the columns we care about
    cols_to_keep = ["title", "description", "main_category", "average_rating", "rating_number"]
    ds = ds.remove_columns([c for c in ds.column_names if c not in cols_to_keep])

    # Add product quality score column
    def add_quality_score(batch):
        return {
            "product_quality_score": [
                product_quality_score(r, n)
                for r, n in zip(batch["average_rating"], batch["rating_number"])
            ]
        }
    ds = ds.map(add_quality_score, batched=True)

    # Only keep rows with valid values
    def is_valid(v):
        """
        Must have valid values in the row. Will be used for filtering.
        """
        if v is None:
            return False
        if isinstance(v, str):
            if v.strip() == "":
              return False
        return True

    def keep_row(row):
        """
        Keep only the columns with valid data
        """
        if is_valid(row.get("title")) and \
          is_valid(row.get("description")) and \
          is_valid(row.get("main_category")) and \
          is_valid(row.get("average_rating")) and \
          is_valid(row.get("rating_number")):
            return True
        return False

    ds = ds.filter(keep_row)

    return ds.to_pandas()

def load_near_field_from_huggingface():
    """
    Load the near-field dataset from HuggingFace.
    """
    ds = load_dataset(REPO_ID_NEAR_FIELD, split="train")
    return ds.to_pandas()

dataset_near_field = load_near_field_from_huggingface()

def load_far_field_from_huggingface():
    """
    Load the far-field dataset from HuggingFace.
    """
    ds = load_dataset(REPO_ID_FAR_FIELD, split="train")
    return ds.to_pandas()

dataset_far_field = load_far_field_from_huggingface()

def product_score(product_quality_score: float, fuzzy_score: float):
    """
    Combine product score and fuzzy score into a single score.
    """
    return math.sqrt(product_quality_score * fuzzy_score)

def query_near_field(input: str, top_k: int=1):
    """
    Return top_k fuzzy matches for query against dataset titles as a pandas DataFrame.
    Always returns exactly top_k rows (if available).
    """
    if top_k <= 0:
        raise ValueError

    n = len(dataset_near_field)
    if top_k > n:
        print(f"Warning: top_k ({top_k}) is greater than the number of examples in the near-field dataset ({n}). Returning all examples.")
        return dataset_near_field.reset_index(drop=True)

    matches = process.extract(
        input,
        dataset_near_field["title"].fillna("").astype(str).tolist(),
        scorer=fuzz.token_set_ratio,
        limit=n
    )

    rows = []
    for _text, fuzzy_score, idx in matches:
        row = dataset_near_field.iloc[idx].to_dict()  # pandas way
        row["data_source"] = "near_field"
        row["fuzzy_score"] = fuzzy_score
        product_quality_score = row.get("product_quality_score")
        row["score"] = product_score(product_quality_score, fuzzy_score)
        rows.append(row)

    return (
        pd.DataFrame(rows)
        .sort_values("score", ascending=False)
        .head(top_k)
        .reset_index(drop=True)
    )

def query_far_field(input: str, top_k: int):
    """
    Return top_k random elements from the far_field dataset as a pandas DataFrame.
    The input string is ignored.
    """
    if top_k < 0:
        raise ValueError

    n = len(dataset_far_field)
    if top_k > n:
        print(f"Warning: top_k ({top_k}) is greater than the number of examples in the far-field dataset ({n}). Returning all examples.")
        return dataset_far_field.reset_index(drop=True)

    # Sample random rows without replacement
    sampled = dataset_far_field.sample(n=top_k, random_state=None).reset_index(drop=True)

    # Add the rest
    sampled["fuzzy_score"] = [
        fuzz.token_set_ratio(str(t) if pd.notna(t) else "", input)
        for t in sampled.get("title", "")
    ]
    product_quality_scores = sampled.get("product_quality_score")
    fuzzy_scores = sampled["fuzzy_score"]
    sampled["score"] = [product_score(a, b) for a, b in zip(product_quality_scores, fuzzy_scores)]
    sampled["data_source"] = "far_field"

    return sampled

def split_near_and_far_fields(total_examples: int, near_far_ratio: float = 0.5):
    """
    Split the examples between near and far field.
    The ratio represents the examples that will be in the near field to total (near + far).
    """
    ratio = near_far_ratio
    # Validate ratio
    if ratio < 0 or ratio > 1:
      raise ValueError("Ratio must be between 0 and 1")
    if total_examples < 2:
      raise ValueError("Total examples must be at least 2")

    near_field_examples = int(total_examples * ratio)
    far_field_examples = total_examples - near_field_examples

    return near_field_examples, far_field_examples

def query(input: str, total_examples: int, near_far_ratio: float = 0.5):
    near_field_examples, far_field_examples = split_near_and_far_fields(total_examples, near_far_ratio)
    far_field_result = query_far_field(input, far_field_examples)
    #print(far_field_result.head())
    near_field_result = query_near_field(input, near_field_examples)
    #print(near_field_result.head())
    result = pd.concat([near_field_result, far_field_result], ignore_index=True)
    return result

# Example
print("Example: ", query("water bottle", total_examples=4, near_far_ratio=0.5))

# Load base + adapter
def lora_load():
    model_name = "distilbert-base-uncased"      # same base you trained on
    
    tokenizer  = AutoTokenizer.from_pretrained(REPO_ID_LORA_GREEN_PATENTS)  # , token=token)
    base_model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)  # , token=token)
    model      = PeftModel.from_pretrained(base_model, REPO_ID_LORA_GREEN_PATENTS)  # , token=token)
    
    clf = pipeline("text-classification", model=model, tokenizer=tokenizer)
    
    # Examples of patents and products  (fixed commas)
    texts = [
      "A biodegradable plastic composition derived from renewable corn starch.",
      "A new synthetic polymer with enhanced tensile strength.",
      "Refreshing Taste: Every bottle of Pure Life Water is enhanced with minerals for a crisp taste that makes drinking water delicious. 12 pack of 16.9 fl oz water bottles.",
      "This 18/8 stainless steel water bottle is designed to last a lifetime. Plastic free & Eco friendly water bottles are a healthier option for you & the planet! However, Water in stainless steel tastes different than plastic, make sure your taste buds are ready for this healthy switch"
    ]
    print(clf(texts))
    return clf

clf = lora_load()

ex_waterbottle_text = [
    "A single use case made with fossil fuels and gasoline.",
    "An eco-friendly, sustainable bottle made with biodegradable plastic."
]
print(clf(ex_waterbottle_text))

def sustainability_filter(input: str, total_examples: int, near_far_ratio: float = 0.5):
  initial_products = query(input, total_examples, near_far_ratio)
  filtered_products = clf(initial_products['description']) # 1 for green patents, 0 otherwise
  sustainable_products = filtered_products.filter(lambda x: x['label'] == 'LABEL_1')
  return sustainable_products

# 👇 Your system prompt (can be empty)
SYSTEM_PROMPT = """
You are a product analyst. You'll receive product description as input, and extract some product functionality and some product values. Each functionality and value should be keywords only.
Product functionality refers to what the product does: its features, technical capabilities, and performance characteristics. It answers the question: “What can this product do?”
Product value refers to the benefit the customer gains from using the product: how it improves their life, solves their problem, or helps them achieve goals. It answers the question: “Why does this matter to the customer?”
Do **not** duplicate an item in both lists. Keep **functionalities** as concrete features. Keep **values** as clear user benefits.
Your Output is a dictionary. Here is the format:
# Your Input:
  <product_description>
# Your Output:
{
  "values": [
    <value1>,
    <value2>,
    ...
  ],
  "functionalities": [
    <function1>,
    <function2>,
    ...
  ]
}
Don't return anything out of the output format.
"""

@dataclass
class LLMConfig:
    model_id: str                       # e.g. "Qwen/Qwen2.5-1.5B-Instruct" or "Qwen/Qwen2.5-3B-Instruct"
    system_prompt: str = ""             # optional system prompt
    max_new_tokens: int = 256
    temperature: float = 0.2
    top_p: float = 0.9
    repetition_penalty: float = 1.05
    use_4bit: bool = True               # good default for Colab VRAM

def create_llm(
    *,
    model_id: str,
    max_new_tokens: int = 256,
    temperature: float = 0.2,
    top_p: float = 0.9,
    repetition_penalty: float = 1.05,
    use_4bit: bool = True
) -> Callable[[str], str]:
    """
    Load an off-the-shelf chat LLM and return a callable llm(prompt) -> str.
    Pass ONLY the model parameters you want. No size mapping. No llama_cpp.
    """

    cfg = LLMConfig(
        model_id=model_id,
        system_prompt=SYSTEM_PROMPT,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        use_4bit=use_4bit,
    )

    has_cuda = torch.cuda.is_available()
    qconfig: Optional[BitsAndBytesConfig] = None
    if has_cuda and cfg.use_4bit:
        qconfig = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

    tokenizer = AutoTokenizer.from_pretrained(cfg.model_id, use_fast=True)
    model = AutoModelForCausalLM.from_pretrained(
        cfg.model_id,
        device_map="auto",
        torch_dtype=torch.bfloat16 if has_cuda else torch.float32,
        quantization_config=qconfig,
    ).eval()

    def _format_messages(user_text: str) -> str:
        msgs = []
        if cfg.system_prompt:
            msgs.append({"role": "system", "content": cfg.system_prompt})
        msgs.append({"role": "user", "content": user_text})

        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
            return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)

        # Fallback if no chat template is present
        sys = f"System: {cfg.system_prompt}\n\n" if cfg.system_prompt else ""
        return f"{sys}User: {user_text}\nAssistant:"

    @torch.inference_mode()
    def llm(prompt: str,
            max_new_tokens: int = None,
            temperature: float = None,
            top_p: float = None,
            repetition_penalty: float = None) -> str:

        text = _format_messages(prompt)
        inputs = tokenizer(text, return_tensors="pt").to(model.device)
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens or cfg.max_new_tokens,
            do_sample=(temperature or cfg.temperature) > 0.0,
            temperature=temperature or cfg.temperature,
            top_p=top_p or cfg.top_p,
            repetition_penalty=repetition_penalty or cfg.repetition_penalty,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        gen = out[0][inputs["input_ids"].shape[-1]:]
        return tokenizer.decode(gen, skip_special_tokens=True).strip()

    print(f"Loaded: {cfg.model_id} | 4-bit: {bool(qconfig)} | Device: {model.device}")
    return llm

def response_to_triplets(product_title, response: str):
    data = json.loads(response)
    #print(data)

    triples_list = []

    for value in data["values"]:
        triples_list.append(f"({product_title}, HAS_VALUE, {value})")

    for func in data["functionalities"]:
        triples_list.append(f"({product_title}, HAS_FUNCTIONALITY, {func})")

    #print(triples_list)
    return triples_list

    llm = create_llm(
    model_id="Qwen/Qwen2.5-1.5B-Instruct",
    max_new_tokens=200,
    temperature=0.2,
    top_p=0.9,
    repetition_penalty=1.05,
    use_4bit=True,   # set False if you have lots of VRAM
)

# Example
if False: # Change to true to check the example
  title = """
  Surge Protector Power Strip - HANYCONY 8 Outlets 4 USB (2 USB C) Charging Ports, Multi Plug Outlet Extender, 5Ft Braided Extension Cord, Flat Plug Wall Mount Desk Charging Station for Home Office ETL
  """
  description = """
  3-side design power strip surge protector with 8AC widely outlets and 4 USB (2 USB C) charging ports can power up to 12 devices simultaneously. That makes it easier to make the plugs not covering any outlet, and the 2.2 inchces widely spced in between outlets, larger than standard socket, fit big adapters without blocking each other. The compact design saves more space, suitable for the home, office, and college dorm room essentials
  """
  product_description = f"{title}\n{description}"
  response = llm(product_description)
  print("Example: \n", response)
  triplets_list = response_to_triplets(title, response)
  print("Example Triplets: \n", triplets_list)

def main(input: str):
  all_triplets_list = []
  '''
  sustainable_results = sustainability_filter(input, total_examples=10, near_far_ratio=0.5)
  for i, product in sustainable_results.iterrows():
      product_title = product["title"]
      response = llm(product_title)
      triplets_list = response_to_triplets(product_title, response)
      for triplet in triplets_list:
        all_triplets_list.append(triplet)
  '''
  all_triplets_list = [
    '(Zojirushi Stainless Steel Mug, HAS_VALUE, temperature regulation)',
    '(Zojirushi Stainless Steel Mug, HAS_VALUE, ease of use)',
    '(Zojirushi Stainless Steel Mug, HAS_VALUE, portability)'
    '(Zojirushi Stainless Steel Mug, HAS_FUNCTIONALITY, vacuum insulation)',
    '(Zojirushi Stainless Steel Mug, HAS_FUNCTIONALITY, durability)'
  ]
  return all_triplets_list

def create_graph_from_triplets(triplets):
    G = nx.DiGraph()
    for triplet in triplets:
        line = str(triplet).strip()
        if not line:
            continue
        # Try comma-delimited with max 2 splits
        parts = [p.strip(" ()") for p in line.split(",", 2)]
        if len(parts) != 3:
            # Fallback: pipe-delimited
            parts = [p.strip(" ()") for p in line.split("|")]
            if len(parts) != 3:
                continue  # malformed, skip
        subject, predicate, obj = parts
        if subject and predicate and obj:
            G.add_edge(subject, obj, label=predicate)
    return G

def nx_to_pyvis(networkx_graph):
    pyvis_graph = Network(notebook=True, cdn_resources='remote')
    for node in networkx_graph.nodes():
        pyvis_graph.add_node(node)
    for edge in networkx_graph.edges(data=True):
        lbl = edge[2].get("label", "")   # ✅ safe access
        pyvis_graph.add_edge(edge[0], edge[1], label=lbl, title=lbl)
    return pyvis_graph

def generateGraph(triples_list):
    triplets = [t.strip() for t in triples_list if t.strip()]
    graph = create_graph_from_triplets(triplets)
    pyvis_network = nx_to_pyvis(graph)

    pyvis_network.toggle_hide_edges_on_drag(True)
    pyvis_network.toggle_physics(False)
    pyvis_network.set_edge_smooth('discrete')

    html = pyvis_network.generate_html()
    html = html.replace("'", "\"")

    return f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera;
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms
    allow-scripts allow-same-origin allow-popups
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""

def pipeline(user_text: str):
    try:
        triples = main(user_text) or []  # ✅ guard against None
        # Normalize tuples/lists to "S, R, O" strings (keeps your existing generateGraph)
        triples_list = []
        for t in triples:
            if isinstance(t, (tuple, list)) and len(t) == 3:
                triples_list.append(f"{t[0]}, {t[1]}, {t[2]}")
            else:
                triples_list.append(str(t))
        return generateGraph(triples_list)
    except Exception:
        return "<pre style='white-space: pre-wrap; font-size:12px; color:#b00;'>" + traceback.format_exc() + "</pre>"

demo = gr.Interface(
    fn=pipeline,
    inputs=gr.Textbox(label="Enter your query / text", value="", lines=6),
    outputs=gr.HTML(),
    title="Knowledge Graph",
    allow_flagging="never",
    live=False,   # set True if you want it to recompute on each keystroke
    css="""
        #component-0, #component-1, #component-2 {
            display: flex;
            justify-content: center;
            align-items: center;
            flex-direction: column;
        }
        .gradio-container {
            justify-content: center !important;
            align-items: center !important;
            text-align: center;
        }
        textarea, iframe {
            margin: 0 auto;
            display: block;
        }
    """
)

demo.launch(quiet=True)