from typing import List
from pydantic import validator

from my.config import BaseConf, SingleOrList, dispatch
from my.utils.seed import seed_everything

import numpy as np
from voxnerf.vox import VOXRF_REGISTRY
from voxnerf.pipelines import train


class VoxConfig(BaseConf):
    model_type:                 str = "VoxRF"
    bbox_len:                   float = 1.5
    grid_size:                  SingleOrList(int) = [128, 128, 128]
    step_ratio:                 float = 0.5
    density_shift:              float = -10.
    ray_march_weight_thres:     float = 0.0001
    c:                          int = 3
    blend_bg_texture:           bool = False
    bg_texture_hw:              int = 64

    @validator("grid_size")
    def check_gsize(cls, grid_size):
        if isinstance(grid_size, int):
            return [grid_size, ] * 3
        else:
            assert len(grid_size) == 3
            return grid_size

    def make(self):
        params = self.dict()
        m_type = params.pop("model_type")
        model_fn = VOXRF_REGISTRY.get(m_type)

        radius = params.pop('bbox_len')
        aabb = radius * np.array([
            [-1, -1, -1],
            [1, 1, 1]
        ])
        model = model_fn(aabb=aabb, **params)
        return model


class TrainerConfig(BaseConf):
    model:      VoxConfig = VoxConfig()
    scene:      str = "lego"
    n_epoch:    int = 2
    bs:         int = 4096
    lr:         float = 0.02

    def run(self):
        args = self.dict()
        args.pop("model")

        model = self.model.make()
        train(model, **args)


if __name__ == "__main__":
    seed_everything(0)
    dispatch(TrainerConfig)