import logging
from pathlib import Path

import torch

from ..base import MatrixDomain
from ..registry import MatrixGeneratorRegistry
from .storage import save_domain, write_meta

logger = logging.getLogger(__name__)


class OfflineGenerationRunner:

    def __init__(
        self,
        registry: MatrixGeneratorRegistry,
        output_dir: Path,
        num_matrices_per_domain: int,
        num_context_pairs: int,
        batch_size: int = 1,
        seed: int = 42,
    ):
        self.registry = registry
        self.output_dir = output_dir
        self.num_matrices_per_domain = num_matrices_per_domain
        self.num_context_pairs = num_context_pairs
        self.batch_size = batch_size
        self.seed = seed

    def run(self):
        self.output_dir.mkdir(parents=True, exist_ok=True)

        for domain_idx, domain in enumerate(self.registry.domains):
            torch.manual_seed(self.seed + domain_idx)
            logger.info("Generating %d matrices for %s", self.num_matrices_per_domain, domain.value)

            samples = []
            for i in range(self.num_matrices_per_domain):
                data = self.registry.generate(domain, self.batch_size, self.num_context_pairs)
                samples.append(data)

                if (i + 1) % 100 == 0:
                    logger.info("  %s: %d/%d", domain.value, i + 1, self.num_matrices_per_domain)

            save_domain(self.output_dir, domain, samples)
            logger.info("Saved %s (%d samples)", domain.value, len(samples))

        write_meta(
            self.output_dir,
            self.registry.domains,
            self.num_matrices_per_domain,
            self.num_context_pairs,
            self.seed,
        )
        logger.info("Done. Meta written to %s", self.output_dir / "meta.json")
