Lazy Call and Partial

import laco
import laco.language as L
import torch
import torch.nn as nn
import torch.optim as optim

# --- Plain Python (eager) ---
# The model exists the instant this line runs.
model_eager = nn.Linear(784, 10)
print(type(model_eager))  # <class 'torch.nn.modules.linear.Linear'>

# --- laco (lazy) ---
# L.call returns a DictConfig — a recipe, not a model.
cfg = L.call(nn.Linear)(in_features=784, out_features=10)
print(type(cfg))  # <class 'omegaconf.DictConfig'>

# Construction happens here:
model_lazy = laco.instantiate(cfg)
print(type(model_lazy))  # <class 'torch.nn.modules.linear.Linear'>

cfg = L.call(nn.Linear)(in_features=784, out_features=10)

print("_target_ :", cfg._target_)    # noqa: LACO001
print("_convert_:", cfg._convert_)   # noqa: LACO001
print("in_features :", cfg.in_features)   # noqa: LACO001
print("out_features:", cfg.out_features)  # noqa: LACO001
print()
print("--- laco.dump (YAML) ---")
print(laco.dump(cfg))

try:
    bad_cfg = L.call(nn.Linear)(in_featurez=784, out_features=10)  # typo!
except TypeError as e:
    print("TypeError caught:")
    print(e)

# WRONG: can't do this at config time — model doesn't exist yet!
#
# optimizer = optim.Adam(model.parameters(), lr=1e-3)  # model is not defined here
#
# This is a plain Python problem: Adam needs params that only exist post-construction.
print("(skipped — intentionally broken approach)")

# Build the optimizer config — no model needed yet.
optimizer_cfg = L.partial(optim.Adam)(lr=1e-3, weight_decay=1e-5)

print("--- optimizer config (YAML) ---")
print(laco.dump(optimizer_cfg))
print("_partial_ field:", optimizer_cfg._partial_)  # noqa: LACO001

# Build the model config and instantiate it.
model_cfg = L.call(nn.Linear)(in_features=784, out_features=10)
model = laco.instantiate(model_cfg)  # nn.Linear is created here

# Instantiate the optimizer config → we get a functools.partial, not an Adam yet.
optimizer_factory = laco.instantiate(optimizer_cfg)
print(type(optimizer_factory))  # <class 'functools.partial'>

# Now supply model.parameters() — Adam is created here.
optimizer = optimizer_factory(model.parameters())
print(type(optimizer))          # <class 'torch.optim.adam.Adam'>
print("lr =", optimizer.param_groups[0]["lr"])

pretrained_weights = torch.zeros(10, 5)  # some pre-computed tensor

cfg_just = L.just(pretrained_weights)

print("_target_:", cfg_just._target_)  # noqa: LACO001  → 'laco.ops.identity'
print("value   :", type(cfg_just.value))  # noqa: LACO001  → torch.Tensor
print()

recovered = laco.instantiate(cfg_just)
print("recovered type  :", type(recovered))
print("values identical:", torch.equal(recovered, pretrained_weights))

import omegaconf

# Demonstrate what L.required[int]() actually returns at runtime:
missing_value = L.required[int]()
print("L.required[int]() is  :", missing_value)
print("Is omegaconf.MISSING  :", missing_value is omegaconf.MISSING)

# A config with a required field — modelled after text_classifier.py
cfg_with_required = L.call(nn.Embedding)(
    num_embeddings=L.required[int](),
    embedding_dim=64,
)

# laco.dump walks the tree and resolves interpolations and mandatory values
# as it goes, so dumping a node that still holds a MISSING field would raise
# MissingMandatoryValue. To inspect the *unresolved* recipe, render it with
# OmegaConf.to_yaml(...), which leaves the ??? sentinel in place.
print("--- YAML (note: num_embeddings shows as ???) ---")
print(omegaconf.OmegaConf.to_yaml(cfg_with_required))

# Instantiating without filling in the required field raises immediately:
try:
    laco.instantiate(cfg_with_required)
except Exception as e:
    print(type(e).__name__)  # MissingMandatoryValue
    print(str(e)[:120])

# === linear_regression.py (annotated) ===

import laco.language as L             # the config DSL
from torch import nn, optim

# --- 1. Hyperparameters ---
# @L.params wraps the class so that attribute access returns OmegaConf
# interpolation strings like "${hps.in_features}" instead of the bare int.
# The IDE still sees `int` thanks to @dataclass_transform.
# (Full coverage of @L.params in NB03.)
@L.params
class hps:
    in_features:    int   = 8
    out_features:   int   = 1
    bias:           bool  = True
    learning_rate:  float = 1e-2
    momentum:       float = 0.9

# --- 2. Model config ---
# L.call(nn.Linear) returns a factory.  Calling it with kwargs produces a
# DictConfig.  root=True is a hint for the type checker (makes the IDE see
# the return as DictConfig rather than nn.Linear; useful for the top-level
# exported symbol).
#
# hps.in_features is "${hps.in_features}" at runtime — an interpolation string
# that OmegaConf resolves to 8 only when the hps node is present *as a sibling*
# in the same config tree.
model = L.call(nn.Linear, root=True)(
    in_features  = hps.in_features,
    out_features = hps.out_features,
    bias         = hps.bias,
)

# --- 3. Optimizer config ---
# L.partial because SGD needs model.parameters() — unavailable at config time.
# laco.instantiate(optimizer) gives back functools.partial(SGD, lr=..., momentum=...).
# The training script then calls optimizer_factory(model.parameters()) to get the
# real SGD instance.
optimizer = L.partial(optim.SGD)(
    lr       = hps.learning_rate,
    momentum = hps.momentum,
)

# --- 4. Bundle the tree ---
# In the real file, `__all__ = ["model", "optimizer", "hps"]` exports all three;
# laco.load assembles them into one DictConfig, which is what lets the
# "${hps.*}" interpolations in `model`/`optimizer` resolve against the `hps`
# node.  Here we reproduce that bundling explicitly.  Dumping `model` on its own
# would raise InterpolationKeyError, because its "${hps.*}" references have no
# `hps` sibling to point at.
import omegaconf

cfg = omegaconf.OmegaConf.create(
    {"hps": hps(), "model": model, "optimizer": optimizer}
)

print("=== full config (model + optimizer + hps) ===")
print(laco.dump(cfg))

# Load the real file and instantiate the model:
cfg = laco.load("configs://examples/linear_regression.py#model")
built_model = laco.instantiate(cfg)
print(built_model)  # Linear(in_features=8, out_features=1, bias=True)

# Load with an override — change in_features without editing the file:
cfg_wide = laco.load(
    "configs://examples/linear_regression.py",
    "hps.in_features=32",
    key="model"
)
built_wide = laco.instantiate(cfg_wide)
print(built_wide)   # Linear(in_features=32, out_features=1, bias=True)