import re
from core.conversion import noop_logger

def evaluate(llm, zip, readmes, log_fn=noop_logger):
  log_fn("TITLE", "\nLooking for code to train the model...")
  overall = "No"


  patterns = {
      'tensorflow': [
          r'model\.(fit|compile|train_on_batch)',
          r'tf\.GradientTape'
      ],
      'pytorch': [
          r'model\.(train|forward)',
          r'loss\.backward',
          r'optimizer\.step',
      ]
  }
  files = [file_path for file_path in zip.namelist() if ((file_path.endswith(".py") | file_path.endswith(".ipynb")))]
  for file_path in files:
      code = zip.open(file_path).read().decode("utf-8")
      for framework, regex_list in patterns.items():
          for pattern in regex_list:
              if re.search(pattern, code):
                  log_fn("LOG", f"Found code for training a model in {framework} framework in file: {file_path}")
                  overall = "Yes"

  for readme in readmes:
    if (readme):
        if (("train" in readme)):
            log_fn("LOG", "Found something about training in README file")
            overall = "Yes"
      
  if (overall == "No"):
    log_fn("ERROR", "Found no code for training the model.")
  return overall