import math
import os
import sys
from pathlib import Path

import numpy as np


PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from onescience.utils.YParams import YParams


def set_default_data_env():
    os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data"))


def resolve_path(path_value):
    path = Path(path_value)
    return path if path.is_absolute() else PROJECT_ROOT / path


def load_config():
    set_default_data_env()
    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
    cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir))
    cfg.fake_data.root_dir = str(resolve_path(cfg.fake_data.root_dir))
    return cfg


def square_boundary(points, boundary_points):
    bottom = [(x, 0.0) for x in points]
    top = [(x, 1.0) for x in points]
    left = [(0.0, y) for y in points[1:-1]]
    right = [(1.0, y) for y in points[1:-1]]
    boundary = np.array(bottom + top + left + right, dtype=np.float64)
    if boundary.shape[0] == 0:
        raise ValueError("resolution is too small to build boundary points")
    repeats = int(math.ceil(boundary_points / boundary.shape[0]))
    return np.tile(boundary, (repeats, 1))[:boundary_points]


def make_beno_arrays(cfg):
    data_cfg = cfg.datapipe.data
    fake_cfg = cfg.fake_data
    total_samples = int(data_cfg.ntrain) + int(data_cfg.ntest)
    resolution = int(data_cfg.resolution)
    boundary_points = int(fake_cfg.boundary_points)
    rng = np.random.default_rng(int(fake_cfg.seed))

    axis = np.linspace(0.0, 1.0, resolution, dtype=np.float64)
    yy, xx = np.meshgrid(axis, axis, indexing="ij")
    coords = np.stack([xx.reshape(-1), yy.reshape(-1)], axis=-1)

    cell_state = np.zeros((resolution, resolution), dtype=np.float64)
    cell_state[0, :] = 1.0
    cell_state[-1, :] = 1.0
    cell_state[:, 0] = 1.0
    cell_state[:, -1] = 1.0
    cell_state = cell_state.reshape(-1)

    boundary_xy = square_boundary(axis, boundary_points)

    rhs = np.zeros((total_samples, resolution * resolution, 4), dtype=np.float64)
    sol = np.zeros((total_samples, resolution * resolution, 1), dtype=np.float64)
    bc = np.zeros((total_samples, boundary_points, 4), dtype=np.float64)

    for sample_idx in range(total_samples):
        phase = 0.25 * sample_idx
        amplitude = 1.0 + 0.1 * sample_idx
        forcing = (
            amplitude * np.sin(np.pi * xx + phase) * np.sin(np.pi * yy)
            + 0.15 * np.cos(2.0 * np.pi * xx) * np.sin(np.pi * yy + phase)
        )
        solution = 0.25 * forcing + 0.05 * (xx + yy)
        if fake_cfg.noise_std:
            solution += float(fake_cfg.noise_std) * rng.standard_normal(solution.shape)

        boundary_value = (
            0.05 * sample_idx
            + 0.1 * boundary_xy[:, 0]
            - 0.08 * boundary_xy[:, 1]
        )

        rhs[sample_idx, :, 0:2] = coords
        rhs[sample_idx, :, 2] = forcing.reshape(-1)
        rhs[sample_idx, :, 3] = cell_state
        sol[sample_idx, :, 0] = solution.reshape(-1)
        bc[sample_idx, :, 0:2] = boundary_xy
        bc[sample_idx, :, 2] = boundary_value

    return rhs, sol, bc


def clear_matching_cache(cfg):
    cache_dir = Path(cfg.datapipe.source.cache_dir)
    prefix = cfg.datapipe.source.file_prefix
    for split, count in (("train", cfg.datapipe.data.ntrain), ("test", cfg.datapipe.data.ntest)):
        cache_file = cache_dir / f"cached_{prefix}_{split}_{count}.pt"
        if cache_file.exists():
            cache_file.unlink()


def main():
    cfg = load_config()
    data_dir = Path(cfg.datapipe.source.data_dir)
    data_dir.mkdir(parents=True, exist_ok=True)

    rhs, sol, bc = make_beno_arrays(cfg)
    prefix = cfg.datapipe.source.file_prefix
    np.save(data_dir / f"RHS_{prefix}_all.npy", rhs)
    np.save(data_dir / f"SOL_{prefix}_all.npy", sol)
    np.save(data_dir / f"BC_{prefix}_all.npy", bc)
    clear_matching_cache(cfg)

    print(f"Fake BENO data written to {data_dir}")
    print(f"RHS shape={rhs.shape}, SOL shape={sol.shape}, BC shape={bc.shape}")


if __name__ == "__main__":
    main()