"""Sample containers and batching for the multiscale model."""

from __future__ import annotations

from dataclasses import dataclass
from typing import List

import torch
from torch import Tensor

from .graph_generator import ConcreteGraph, MortarMicrograph


@dataclass
class ConcreteSample:
    graph: ConcreteGraph
    micrograph: MortarMicrograph
    mix_itz: Tensor
    concrete_target: Tensor
    # NaN marks an unobserved target.
    concrete_aux_target: Tensor
    mortar_target: Tensor
    itz_target_graph: Tensor
    paste_target: Tensor
    sand_volume_fraction: Tensor
    tabular_mix: Tensor


@dataclass
class MultiscaleBatch:
    x: Tensor
    x_mask: Tensor
    edge_index: Tensor
    edge_attr: Tensor
    edge_attr_mask: Tensor
    global_features: Tensor
    global_mask: Tensor
    batch: Tensor
    pos: Tensor
    node_type: Tensor
    edge_type: Tensor
    m_x: Tensor
    m_mask: Tensor
    m_edge_index: Tensor
    m_edge_attr: Tensor
    m_node_type: Tensor
    m_batch: Tensor
    mortar_global: Tensor
    mortar_global_mask: Tensor
    mix_itz: Tensor
    concrete_target: Tensor
    mortar_target: Tensor
    itz_target_graph: Tensor
    paste_target: Tensor
    sand_volume_fraction: Tensor
    tabular_mix: Tensor
    concrete_aux_target: Tensor

    @property
    def num_graphs(self) -> int:
        return self.concrete_target.size(0)

    def to(self, device: torch.device) -> "MultiscaleBatch":
        return MultiscaleBatch(
            **{
                name: value.to(device)
                for name, value in vars(self).items()
            }
        )


def collate_multiscale(samples: List[ConcreteSample]) -> MultiscaleBatch:
    xs, x_masks, e_idx, e_attr, e_attr_masks = [], [], [], [], []
    globals_, global_masks = [], []
    batches, positions = [], []
    node_types, edge_types = [], []
    m_xs, m_masks, m_eidx, m_eattr, m_node_types, m_batches = [], [], [], [], [], []
    mortar_globals, mortar_global_masks = [], []
    mix_itzs, ys = [], []
    mortar_targets, itz_graph_targets, paste_targets = [], [], []
    sand_vfs, tabular_mixes, concrete_aux_targets = [], [], []

    node_offset = 0
    micro_offset = 0
    for gid, sample in enumerate(samples):
        graph: ConcreteGraph = sample.graph
        micro: MortarMicrograph = sample.micrograph
        n_nodes = graph.x.size(0)
        n_micro = micro.x.size(0)

        xs.append(graph.x)
        x_masks.append(graph.x_mask)
        e_idx.append(graph.edge_index + node_offset)
        e_attr.append(graph.edge_attr)
        e_attr_masks.append(graph.edge_attr_mask)
        globals_.append(graph.global_features)
        global_masks.append(graph.global_mask)
        batches.append(torch.full((n_nodes,), gid, dtype=torch.long))
        positions.append(graph.pos)
        node_types.append(graph.node_type)
        edge_types.append(graph.edge_type)

        m_xs.append(micro.x)
        m_masks.append(micro.x_mask)
        m_eidx.append(micro.edge_index + micro_offset)
        m_eattr.append(micro.edge_attr)
        m_node_types.append(micro.node_type)
        m_batches.append(torch.full((n_micro,), gid, dtype=torch.long))
        mortar_globals.append(micro.global_features)
        mortar_global_masks.append(micro.global_mask)

        mix_itzs.append(sample.mix_itz)
        ys.append(sample.concrete_target)
        mortar_targets.append(sample.mortar_target)
        itz_graph_targets.append(sample.itz_target_graph)
        paste_targets.append(sample.paste_target)
        sand_vfs.append(sample.sand_volume_fraction)
        concrete_aux_targets.append(sample.concrete_aux_target)
        tabular_mixes.append(sample.tabular_mix)

        node_offset += n_nodes
        micro_offset += n_micro

    return MultiscaleBatch(
        x=torch.cat(xs, dim=0),
        x_mask=torch.cat(x_masks, dim=0),
        edge_index=torch.cat(e_idx, dim=1),
        edge_attr=torch.cat(e_attr, dim=0),
        edge_attr_mask=torch.cat(e_attr_masks, dim=0),
        global_features=torch.stack(globals_, dim=0),
        global_mask=torch.stack(global_masks, dim=0),
        batch=torch.cat(batches, dim=0),
        pos=torch.cat(positions, dim=0),
        node_type=torch.cat(node_types, dim=0),
        edge_type=torch.cat(edge_types, dim=0),
        m_x=torch.cat(m_xs, dim=0),
        m_mask=torch.cat(m_masks, dim=0),
        m_edge_index=torch.cat(m_eidx, dim=1),
        m_edge_attr=torch.cat(m_eattr, dim=0),
        m_node_type=torch.cat(m_node_types, dim=0),
        m_batch=torch.cat(m_batches, dim=0),
        mortar_global=torch.stack(mortar_globals, dim=0),
        mortar_global_mask=torch.stack(mortar_global_masks, dim=0),
        mix_itz=torch.stack(mix_itzs, dim=0),
        concrete_target=torch.stack(ys, dim=0),
        mortar_target=torch.stack(mortar_targets, dim=0),
        itz_target_graph=torch.stack(itz_graph_targets, dim=0),
        paste_target=torch.stack(paste_targets, dim=0),
        sand_volume_fraction=torch.stack(sand_vfs, dim=0),
        tabular_mix=torch.stack(tabular_mixes, dim=0),
        concrete_aux_target=torch.stack(concrete_aux_targets, dim=0),
    )
