from datetime import datetime
from pathlib import Path

import yaml


# Define the model
class MyModel:
    def __init__(
        self,
        n_layers: int,
        n_heads: int,
        dropout: float,
        outdir: str | Path,
    ):
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.dropout = dropout
        self.outdir = outdir

    @classmethod
    def from_config(cls, config: Path | str):
        with open(config, "r") as f:
            hparams = yaml.safe_load(f)
        return cls(**hparams)

    def load_checkpoint(self, checkpoint: Path | str):
        print(f"Loading checkpoint from {checkpoint}...")

    def train(self):
        print("Training...")
        print("Saving checkpoint...")
        with open(Path(self.outdir) / "checkpoint.pt", "w") as f:
            f.write("Hello, world!")

    def evaluate(self):
        print("Evaluating...")


# Create the output directory
now = datetime.now().strftime("%Y%m%d_%H%M%S")
outdir = Path(f"{now}_runs")
outdir.mkdir(parents=True, exist_ok=True)

# Instantiate the model
hparams = {
    "n_layers": 3,
    "n_heads": 8,
    "dropout": 0.1,
    "outdir": str(outdir),
}

model = MyModel(**hparams)

# Train the model
model.train()

# Save the hyperparameters in the same directory
with open(outdir / "hparams.yaml", "w") as f:
    yaml.dump(hparams, f)

# Evaluate the model
new_model = MyModel.from_config(outdir / "hparams.yaml")
new_model.load_checkpoint(outdir / "checkpoint.pt")
new_model.evaluate()