# Path: encore/model_factory.py
# class ModelFactory:
#     __models = []
# 
#     @staticmethod
#     def list(config = {}):
#         def desc(x):
#             return {"code": x.model_code, "name": x.model_name, "description": x.model_desc,
#                 "depends": getattr(x, "depends",  []),
#                 "filters": getattr(x, "filters",  []),
#                 "response_class": getattr(x, "response_class", "*")}
#         show_models = {x: 1 for x in config.get("AVAIL_MODELS", [])}
#         def can_show(x):
#             return len(show_models)==0 or x.model_code in show_models
#         return [ desc(m) for m in ModelFactory.__models if can_show(m)]
# 
#     @staticmethod
#     def register(m):
#         if not hasattr(m, "model_code"):
#             raise ValueError("Missing required attribite 'model_code'")
#         ModelFactory.__models.append(m)
# 
#     @staticmethod
#     def get(model_code, working_directory, config):
#         for m in ModelFactory.__models:
#             if m.model_code == model_code:
#                 return m(working_directory, config)
#         raise ValueError("Unrecognized model type: {}".format(model_code))
# 
#     @staticmethod
#     def get_for_model_spec(model_spec, working_directory, config):
#         model_code = model_spec.get("type", None)
#         if model_code is None:
#             raise ValueError("Type not found for model")
#         return ModelFactory.get(model_code, working_directory, config)


import json
import os
import subprocess
import pwd
from .model_factory import ModelFactory
    def get_batch_headers(self, model_plan):
        mem_per_cpu = model_plan.get("mem_per_cpu", self.config.get("JOB_MEM_PER_CPU", 6500))
        cores_per_job = model_plan.get("cores_per_job", self.config.get("JOB_CORES_PER_TASK", 56))
        
        sbatch_headers = ["#!/bin/bash"]

        if "SLURM_ACCOUNT" in self.config:
            sbatch_headers.append(
                "#SBATCH --account={}".format(self.config.get("SLURM_ACCOUNT")))

        sbatch_headers.extend((
            "#SBATCH --partition={}".format(self.config.get("QUEUE_PARTITION", "encore")),
            "#SBATCH --job-name=gasp_{}".format(self.job_id),
            "#SBATCH --mem-per-cpu={}".format(mem_per_cpu),
            "#SBATCH --chdir={}".format(self.job_directory),
            "#SBATCH --cpus-per-task={}".format(cores_per_job),
            "#SBATCH --time={}".format(self.config.get("JOB_TIME", "14-0")),
            "#SBATCH --nodes=1",
            "export OPENBLAS_NUM_THREADS=1"))

        return sbatch_headers

    def write_batch_script(self, batch_script_path, model_plan):
        with open(batch_script_path, "w") as f:
            f.write("\n".join(self.get_batch_headers(model_plan)))
            f.write("\n\n")
            f.write("\n".join(model_plan["commands"]))
            f.write("\n")

    def submit_job(self, model_spec):

# Path: encore/model_factory.py
# class ModelFactory:
#     __models = []
# 
#     @staticmethod
#     def list(config = {}):
#         def desc(x):
#             return {"code": x.model_code, "name": x.model_name, "description": x.model_desc,
#                 "depends": getattr(x, "depends",  []),
#                 "filters": getattr(x, "filters",  []),
#                 "response_class": getattr(x, "response_class", "*")}
#         show_models = {x: 1 for x in config.get("AVAIL_MODELS", [])}
#         def can_show(x):
#             return len(show_models)==0 or x.model_code in show_models
#         return [ desc(m) for m in ModelFactory.__models if can_show(m)]
# 
#     @staticmethod
#     def register(m):
#         if not hasattr(m, "model_code"):
#             raise ValueError("Missing required attribite 'model_code'")
#         ModelFactory.__models.append(m)
# 
#     @staticmethod
#     def get(model_code, working_directory, config):
#         for m in ModelFactory.__models:
#             if m.model_code == model_code:
#                 return m(working_directory, config)
#         raise ValueError("Unrecognized model type: {}".format(model_code))
# 
#     @staticmethod
#     def get_for_model_spec(model_spec, working_directory, config):
#         model_code = model_spec.get("type", None)
#         if model_code is None:
#             raise ValueError("Type not found for model")
#         return ModelFactory.get(model_code, working_directory, config)


# Path: encore/slurm_queue.py
import json
import os
import subprocess
import pwd
from .model_factory import ModelFactory

    def get_batch_headers(self, model_plan):
        mem_per_cpu = model_plan.get("mem_per_cpu", self.config.get("JOB_MEM_PER_CPU", 6500))
        cores_per_job = model_plan.get("cores_per_job", self.config.get("JOB_CORES_PER_TASK", 56))
        
        sbatch_headers = ["#!/bin/bash"]

        if "SLURM_ACCOUNT" in self.config:
            sbatch_headers.append(
                "#SBATCH --account={}".format(self.config.get("SLURM_ACCOUNT")))

        sbatch_headers.extend((
            "#SBATCH --partition={}".format(self.config.get("QUEUE_PARTITION", "encore")),
            "#SBATCH --job-name=gasp_{}".format(self.job_id),
            "#SBATCH --mem-per-cpu={}".format(mem_per_cpu),
            "#SBATCH --chdir={}".format(self.job_directory),
            "#SBATCH --cpus-per-task={}".format(cores_per_job),
            "#SBATCH --time={}".format(self.config.get("JOB_TIME", "14-0")),
            "#SBATCH --nodes=1",
            "export OPENBLAS_NUM_THREADS=1"))

        return sbatch_headers

    def write_batch_script(self, batch_script_path, model_plan):
        with open(batch_script_path, "w") as f:
            f.write("\n".join(self.get_batch_headers(model_plan)))
            f.write("\n\n")
            f.write("\n".join(model_plan["commands"]))
            f.write("\n")

    def submit_job(self, model_spec):

        model = ModelFactory.get_for_model_spec(model_spec, self.job_directory, self.config)