from .utils import log
import re

def evaluate(verbose, llm, zip, readme):
  log(verbose, "LOG", "\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(verbose, "LOG", f"Found code for training a model in {framework} framework in file: {file_path}")
                  overall = "Yes"


  if (readme):
      if (("train" in readme)):
          log(verbose, "LOG", "Found something about training in README file")
          overall = "Yes"

  return overall