import re
from core.conversion import noop_logger

def evaluate(llm, zip, readmes, log_fn=noop_logger):
  log_fn("TITLE", "\nLooking for pre-trained model weights...")
  overall = "No"
  files = [file_path for file_path in zip.namelist() if ((file_path.endswith(".h5") | file_path.endswith(".pth") | file_path.endswith(".torch") | file_path.endswith(".pt") | file_path.endswith(".tar.gz") | file_path.endswith("checkpoint.pt") | ("weights" in file_path) | file_path.endswith("ckpt")))]
  if (len(files) > 0):
      log_fn("LOG", f"Found model weights: {files}")
      overall = "Yes"
      return overall

  for readme in readmes:
    if (readme):

        url_pattern = r'(https?://[^\s]+)'
        urls = re.findall(url_pattern, readme)
        if (len([url for url in urls if "pth" in url]) > 0):
            log_fn("LOG", "Found a link to pre-trained weights in readme")
            overall = "Yes"
            return overall

        readme_lines = readme.split("\n")
        if (len([row for row in readme_lines if ((len(re.findall("pretrained", row, re.IGNORECASE)) > 0) & (len(re.findall("http", row, re.IGNORECASE)) > 0))]) > 0):
            log_fn("LOG", "Found a link for 'pretrained' something in readme")
            overall = "Yes"
            return overall

        if (len([row for row in readme_lines if ((len(re.findall("pre-trained", row, re.IGNORECASE)) > 0) & (len(re.findall("http", row, re.IGNORECASE)) > 0))]) > 0):
            log_fn("LOG", "Found a link for 'pre-trained' something in readme")
            overall = "Yes"
            return overall

        if (len([row for row in readme_lines if ((len(re.findall("weight", row, re.IGNORECASE)) > 0) & (len(re.findall("http", row, re.IGNORECASE)) > 0))]) > 0):
            log_fn("LOG", "Found a link for 'weight' something in readme")
            overall = "Yes"
            return overall

        if (len([row for row in readme_lines if ((len(re.findall("download", row, re.IGNORECASE)) > 0) & (len(re.findall("model", row, re.IGNORECASE)) > 0) & (len(re.findall("http", row, re.IGNORECASE)) > 0))]) > 0):
            log_fn("LOG", "Found a link for 'model' something in readme")
            overall = "Yes"
            return overall

        if (llm):
            prompt = f"{readme}\nQ: Does this text contain a download link for the model pre-trained weights?"
            ans = llm.predict("STRICT", prompt)
            if (("Yes" in ans) & ("No" not in ans)):
                log_fn("LOG", "The LLM found signs for accessing the pre-trained weights from the readme")
                overall = "Yes"
                return overall

  log_fn("ERROR", "Found no pre-trained model weights.")
  return overall