from pysbmy.density import mesh_to_mesh
from pysbmy.field import Field, read_field
import numpy as np
import os


def field_to_field(
        input_file:str,
        output_file:str,
        output_size:int|tuple[int,int,int]|list[int],
        output_L:float|tuple[float,float,float]|list[float],
        output_dpm:float|tuple[float,float,float]|list[float],
        output_corner:tuple[float,float,float]|list[float],
        boundary_conditions:int,
        ):
    

    ### Make sure all inputs are valid
    if output_L is not None and output_dpm is not None:
        raise ValueError("Either output_L or output_dpm can be specified, not both.")

    if isinstance(output_size, int):
        output_size = (output_size, output_size, output_size) # N0, N1, N2
    elif len(output_size) == 1:
        output_size = (output_size[0], output_size[0], output_size[0])

    if output_L is not None:
        if isinstance(output_L, float):
            output_L = (output_L, output_L, output_L)
        elif len(output_L) == 1:
            output_L = (output_L[0], output_L[0], output_L[0])
    
    if output_dpm is None:
        if output_L is None:
            output_dpm = (-1, -1, -1)
        else:
            output_dpm = (output_L[0] / output_size[0], output_L[1] / output_size[1], output_L[2] / output_size[2])
    elif isinstance(output_dpm, float):
        output_dpm = (output_dpm, output_dpm, output_dpm) # d0, d1, d2
    elif len(output_dpm) == 1:
        output_dpm = (output_dpm[0], output_dpm[0], output_dpm[0])

    if not os.path.exists(input_file):
        raise FileNotFoundError(f"Input file {input_file} does not exist.")
    

    
    # Read the input field
    print(f"Reading input field from {input_file}")
    input_field = read_field(input_file)

    if input_field.rank != 1 or input_field.data.ndim != 3:
        raise NotImplementedError("Only 3D scalar fields are supported for now.")

    L0 = input_field.L0
    L1 = input_field.L1
    L2 = input_field.L2
    N0 = input_field.N0
    N1 = input_field.N1
    N2 = input_field.N2
    corner0 = input_field.corner0
    corner1 = input_field.corner1
    corner2 = input_field.corner2

    d0_in = L0 / N0
    d1_in = L1 / N1
    d2_in = L2 / N2

    d0_out = output_dpm[0] if output_dpm[0] > 0 else d0_in
    d1_out = output_dpm[1] if output_dpm[1] > 0 else d1_in
    d2_out = output_dpm[2] if output_dpm[2] > 0 else d2_in

    offset_out_x = (output_corner[0] - corner0)
    offset_out_y = (output_corner[1] - corner1)
    offset_out_z = (output_corner[2] - corner2)
    
    print("-----------------------------------------")
    print(f"Input field size:    {N0} x {N1} x {N2}")
    print(f"Output field size:   {output_size[0]} x {output_size[1]} x {output_size[2]}")
    print(f"Input field dpm:     {d0_in:.3f} x {d1_in:.3f} x {d2_in:.3f}")
    print(f"Output field dpm:    {d0_out:.3f} x {d1_out:.3f} x {d2_out:.3f}")
    print(f"Input field corner:  ({corner0:.1f}, {corner1:.1f}, {corner2:.1f})")
    print(f"Output field corner: ({output_corner[0]:.1f}, {output_corner[1]:.1f}, {output_corner[2]:.1f})")
    print(f"Boundary conditions: {'periodic' if boundary_conditions == 1 else 'non-periodic'}")
    print("-----------------------------------------")

    input_grid = input_field.data

    output_grid = np.zeros(output_size, dtype=input_grid.dtype)


    # Mesh to mesh interpolation
    print("Interpolating field...")
    mesh_to_mesh(input_grid, output_grid, d0_in, d1_in, d2_in, d0_out, d1_out, d2_out, offset_out_x, offset_out_y, offset_out_z, boundary_conditions)

    # Create the output field
    output_field = Field(
        L0=output_size[0] * d0_out,
        L1=output_size[1] * d1_out,
        L2=output_size[2] * d2_out,
        corner0=output_corner[0],
        corner1=output_corner[1],
        corner2=output_corner[2],
        rank=input_field.rank,
        N0=output_size[0],
        N1=output_size[1],
        N2=output_size[2],
        time=input_field.time,
        data=output_grid
    )

    # Write the output field
    print(f"Writing output field to {output_file}")
    output_field.write(output_file)
    print("Done.")


def console_main():
    import argparse
    parser = argparse.ArgumentParser(description="Convert a field from one size to another.")
    parser.add_argument("-i","--input_file", type=str, help="Input field file")
    parser.add_argument("-o","--output_file", type=str, help="Output field file")
    parser.add_argument("-N","--output_size", type=int, nargs="+", help="Output field size (N0, N1, N2)")
    parser.add_argument("-L","--output_L", type=float, nargs="+", default=None, help="Output field size (L0, L1, L2)")
    parser.add_argument("-dpm","--output_dpm", type=float, nargs="+", default=None, help="Output field dpm (d0, d1, d2)")
    parser.add_argument("-corner","--output_corner", type=float, nargs=3, default=(0.,0.,0.), help="Output field corner (corner0, corner1, corner2)")
    parser.add_argument("-BC","--boundary_conditions", type=int, default=1, help="Boundary conditions (1: periodic, 3: non-periodic)")
    args = parser.parse_args()

    field_to_field(args.input_file, args.output_file, args.output_size, args.output_L, args.output_dpm, args.output_corner, args.boundary_conditions)


if __name__ == "__main__":
    console_main()