#!/usr/bin/env python3

import os
import sys
import glob
import math
import argparse

import cv2
import numpy as np
import PIL.Image
from PIL import Image

# We allow these file extensions as inputs
SUPPORTED_EXTS = ['.tga', '.png', '.jpg', '.jpeg']

def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert a 6-face cubemap into a reflection-mapped tiny-planet image.'
    )
    parser.add_argument(
        '--input_path', '-i',
        default='.',
        help='Path containing the cubemap face images. Defaults to current folder.'
    )
    parser.add_argument(
        '--output_path', '-o',
        default=None,
        help='Folder for output images. Defaults to the same as input_path.'
    )
    parser.add_argument(
        '--prefix', '-p',
        default=None,
        help='Filename prefix for the cubemap faces, e.g. "mymap" -> "mymapft.tga", etc. '
             'If not provided, the script tries to infer it from the first face it finds.'
    )
    parser.add_argument(
        '--type', '-t',
        default=None,
        help='Output file format extension (png, jpg, tga, etc.). '
             'If not provided, defaults to the extension of the first face found.'
    )
    parser.add_argument(
        '--size', '-s',
        default='2048x1024',
        help='Equirectangular output size, WIDTHxHEIGHT (default=2048x1024).'
    )
    parser.add_argument(
        '--pan',
        type=float,
        default=None,
        help='Rotate around X axis by this many degrees. '
             'By default, if reflection map is on, pan=180, else pan=0.'
    )
    parser.add_argument(
        '--roll',
        type=float,
        default=0.0,
        help='Rotate around Y axis by this many degrees (default=0).'
    )
    parser.add_argument(
        '--tilt',
        type=float,
        default=0.0,
        help='Rotate around Z axis by this many degrees (default=0).'
    )
    parser.add_argument(
        '--flipX',
        action='store_true',
        help='Flip each face horizontally before stitching.'
    )
    parser.add_argument(
        '--flipY',
        action='store_true',
        help='Flip each face vertically before stitching.'
    )
    # Reflection map is ON by default. If --noreflmap is used, we disable it.
    parser.add_argument(
        '--mirror',
        action='store_true',
        help='(Advanced) Force reflection map mode ON (vertically mirror each face). '
             'But reflection map is already ON by default unless --noreflmap is given.'
    )
    parser.add_argument(
        '--noreflmap',
        action='store_true',
        help='Disable reflection map (no vertical mirroring).'
    )

    return parser.parse_args()

def read_image_with_pillow(path):
    """
    Read an image with Pillow (supports TGA better), convert to NumPy BGR.
    """
    im = Image.open(path).convert('RGB')
    np_img = np.array(im)[:, :, ::-1]  # swap R <-> B
    return np_img

def cv_write_tga(filepath, bgr_array):
    # bgr_array is e.g. equirect (OpenCV style)
    # convert to RGB
    rgb_array = cv2.cvtColor(bgr_array, cv2.COLOR_BGR2RGB)
    pil_img = PIL.Image.fromarray(rgb_array)
    pil_img.save(filepath, format='TGA')

def my_imread(path):
    """
    Attempt to read using OpenCV; if extension is .tga or if OpenCV fails,
    fallback to Pillow.
    """
    ext = os.path.splitext(path)[1].lower()
    if ext == '.tga':
        return read_image_with_pillow(path)
    else:
        img = cv2.imread(path, cv2.IMREAD_COLOR)
        if img is None:
            return read_image_with_pillow(path)
        return img

def euler_rotation_matrix(pan_deg, roll_deg, tilt_deg):
    """
    Create a rotation matrix from euler angles:
      pan = rotation around X axis
      roll = rotation around Y axis
      tilt = rotation around Z axis
    """
    rx = math.radians(pan_deg)
    ry = math.radians(roll_deg)
    rz = math.radians(tilt_deg)

    Rx = np.array([
        [1,           0,            0],
        [0,  math.cos(rx), -math.sin(rx)],
        [0,  math.sin(rx),  math.cos(rx)]
    ], dtype=np.float64)

    Ry = np.array([
        [math.cos(ry),  0, math.sin(ry)],
        [          0,   1,          0 ],
        [-math.sin(ry), 0, math.cos(ry)]
    ], dtype=np.float64)

    Rz = np.array([
        [math.cos(rz), -math.sin(rz), 0],
        [math.sin(rz),  math.cos(rz), 0],
        [          0,             0,  1]
    ], dtype=np.float64)

    # final = Rz * Ry * Rx  (apply X then Y then Z)
    return Rz @ Ry @ Rx

#
# By default, we look for these suffixes in the following order:
#    front => "ft", back => "bk", left => "lt", right => "rt", up => "up", down => "dn"
# or possibly something else depending on user’s arrangement.  
# But the user specifically said they are using "ft, bk, lt, rt, up, dn".
# Adjust as needed for your naming scheme.
#
SUFFIX_MAP = {
    'front': 'ft',
    'back':  'bk',
    'left':  'lf',
    'right': 'rt',
    'up':    'up',
    'down':  'dn',
}

def extract_prefix_and_ext(filename, expected_suffix):
    """
    Given a filename (without path) and a suffix like 'ft', 
    try to remove that suffix from the basename. Return (prefix, ext).
    E.g. 'myskyft.tga' with suffix='ft' => prefix='mysky', ext='.tga'.
         'earthbk.png' with suffix='bk' => prefix='earth', ext='.png'.
    If suffix not found in the correct place, we just return 
    (filename_without_ext, extension).
    """
    base, ext = os.path.splitext(filename)
    ext = ext.lower()
    # if base ends with expected_suffix, remove it
    if base.endswith(expected_suffix):
        prefix_candidate = base[: -len(expected_suffix)]
    else:
        prefix_candidate = base
    return prefix_candidate, ext

def load_cubemap_faces(input_path, prefix=None, flip_x=False, flip_y=False, mirror=False):
    """
    Load the 6 faces of a cubemap. The SUFFIX_MAP defines which suffix corresponds to each face key.
    If prefix is not specified, we try to infer it from the first face we find.

    We also detect the extension from the first face found if the user didn't specify.

    Then we apply:
      - Horizontal flip if flip_x
      - Vertical flip if flip_y
      - If mirror is True, we create a taller image by placing the V-flipped copy on top.

    Returns:
      (faces_dict, inferred_prefix, inferred_ext)
      where faces_dict = {
         'front': BGR image,
         'back':  BGR image,
         'left':  BGR image,
         'right': BGR image,
         'up':    BGR image,
         'down':  BGR image
      }
    """
    faces = {}
    inferred_prefix = prefix if prefix else None
    inferred_ext = None

    for face_key, suffix in SUFFIX_MAP.items():
        if prefix:
            pattern = os.path.join(input_path, prefix + suffix)
        else:
            # e.g. "*ft"
            pattern = os.path.join(input_path, '*' + suffix)

        # gather candidates
        candidates = []
        for ext in SUPPORTED_EXTS:
            glob_pattern = pattern + ext
            found = glob.glob(glob_pattern)
            candidates.extend(found)

        if not candidates:
            raise FileNotFoundError(
                f"Could not find a file for face '{face_key}' with suffix '{suffix}' in {input_path}."
            )

        # Use the first matching candidate
        filepath = candidates[0]
        img = my_imread(filepath)
        if img is None:
            raise IOError(f"Failed to read {filepath}")

        # If we haven't inferred prefix/ext yet and prefix is None, parse them
        if inferred_prefix is None or inferred_ext is None:
            fname = os.path.basename(filepath)
            deduced_prefix, deduced_ext = extract_prefix_and_ext(fname, suffix)
            # store them if we haven't set them yet
            if inferred_prefix is None:
                inferred_prefix = deduced_prefix
            if inferred_ext is None:
                # Map .jpeg => .jpg for consistency
                if deduced_ext in ['.jpeg', '.jpg']:
                    deduced_ext = '.jpg'
                inferred_ext = deduced_ext

        # Apply flips
        if flip_x:
            # horizontal flip
            img = cv2.flip(img, 1)
        if flip_y:
            # vertical flip
            img = cv2.flip(img, 0)

        # If reflection map is ON, mirror each face vertically:
        # "mirror above itself" => top half is a vertically flipped copy
        # bottom half is the original.
        if mirror:
            flipped = cv2.flip(img, 0)  # vertical flip
            img = np.concatenate([flipped, img], axis=0)

        faces[face_key] = img

    return faces, inferred_prefix, inferred_ext

def rotate_direction(dx, dy, dz, R):
    v = np.array([dx, dy, dz], dtype=np.float64)
    vr = R @ v
    return vr[0], vr[1], vr[2]

def direction_to_cubemap_uv(dx, dy, dz):
    """
    Standard face orientations:
      +X=right, -X=left, +Y=up, -Y=down, +Z=front, -Z=back.
    The user said they'd rearranged faces so that the final 'left' and 'right' had negative U, etc.
    We'll keep the logic as in the existing script (with the user’s final adjustments).
    """
    absX, absY, absZ = abs(dx), abs(dy), abs(dz)
    major_axis = max(absX, absY, absZ)

    if major_axis < 1e-12:
        # degenerate direction
        return 'front', 0.5, 0.5

    if major_axis == absX:
        if dx > 0:
            # +X face = 'right'
            face = 'right'
            # user used negative for U
            u = -0.5 * (dz / absX) + 0.5
            v = -0.5 * (dy / absX) + 0.5
        else:
            # -X face = 'left'
            face = 'left'
            u = -0.5 * (-dz / absX) + 0.5
            v = -0.5 * ( dy / absX) + 0.5

    elif major_axis == absY:
        if dy > 0:
            # +Y face = 'up'
            face = 'up'
            u = 0.5 * (dx / absY) + 0.5
            v = 0.5 * (dz / absY) + 0.5
        else:
            # -Y face = 'down'
            face = 'down'
            u = 0.5 * (dx / absY) + 0.5
            v = 0.5 * (-dz / absY) + 0.5

    else:  # major_axis == absZ
        if dz > 0:
            # +Z face = 'front'
            face = 'front'
            u = 0.5 * (dx / absZ) + 0.5
            v = -0.5 * (dy / absZ) + 0.5
        else:
            # -Z face = 'back'
            face = 'back'
            u = 0.5 * (-dx / absZ) + 0.5
            v = -0.5 * (dy / absZ) + 0.5

    return face, u, v

def cubemap_to_equirectangular(faces, out_w=2048, out_h=1024, R=None):
    """
    Generate a standard equirectangular from the 6 faces in 'faces'.
    If R is provided, rotate each direction vector before sampling.
    """
    equi = np.zeros((out_h, out_w, 3), dtype=np.uint8)

    for j in range(out_h):
        # latitude phi in [0..pi], top=0, bottom=pi
        phi = math.pi * (1.0 - j / (out_h - 1.0))

        for i in range(out_w):
            # longitude theta in [0..2*pi]
            theta = 2.0 * math.pi * (i / (out_w - 1.0))

            sin_phi = math.sin(phi)
            dx = sin_phi * math.cos(theta)
            dy = math.cos(phi)
            dz = sin_phi * math.sin(theta)

            if R is not None:
                dx, dy, dz = rotate_direction(dx, dy, dz, R)

            face_name, u, v = direction_to_cubemap_uv(dx, dy, dz)
            face_img = faces[face_name]
            fh, fw, _ = face_img.shape

            px = int(u * (fw - 1))
            py = int(v * (fh - 1))
            px = max(0, min(px, fw - 1))
            py = max(0, min(py, fh - 1))

            equi[j, i] = face_img[py, px]

    return equi

def stereographic_tiny_planet(equi, out_size=1024, fov_degrees=360):
    """
    Convert equirectangular into a stereographic (tiny planet).
    We'll do a radial map: r in [0..radius] => lat in [0..pi].
    """
    h, w, _ = equi.shape
    cx = out_size / 2.0
    cy = out_size / 2.0
    radius = out_size / 2.0

    planet = np.zeros((out_size, out_size, 3), dtype=np.uint8)

    for y in range(out_size):
        for x in range(out_size):
            dx = x - cx
            dy = y - cy
            r = math.sqrt(dx*dx + dy*dy)
            alpha = math.atan2(dy, dx)  # [-pi..pi]

            if r < 1e-9:
                # near center => top row middle
                planet[y, x] = equi[0, w // 2]
                continue

            # fraction of radius => [0..1]
            r_frac = r / radius
            # map r_frac => lat in [0..pi]
            phi = r_frac * math.pi
            lon = alpha
            if lon < 0:
                lon += 2.0 * math.pi

            # equirect coords
            equirect_x = (lon / (2.0 * math.pi)) * w
            equirect_y = (1.0 - (phi / math.pi)) * h

            ix = int(np.clip(equirect_x, 0, w-1))
            iy = int(np.clip(equirect_y, 0, h-1))

            planet[y, x] = equi[iy, ix]

    return planet

def main():
    args = parse_args()

    # Decide whether reflection map is on or off:
    # It's ON by default, disabled by --noreflmap
    # If user explicitly used --mirror, that also sets it ON (but we handle the default anyway).
    # We'll finalize it in this order:
    #   If --noreflmap: mirror=False
    #   else: mirror=True
    # (the user can still pass --mirror, but it's already the default, so no difference).
    if args.noreflmap:
        mirror = False
    else:
        mirror = True

    # If output_path not given, default to input_path
    if not args.output_path:
        args.output_path = args.input_path

    # Parse desired equirect size
    try:
        w_str, h_str = args.size.lower().split('x')
        out_w, out_h = int(w_str), int(h_str)
    except:
        print(f"Invalid --size argument: {args.size}. Must be WIDTHxHEIGHT, e.g. 2048x1024.")
        sys.exit(1)

    # 1) Load the 6 cubemap faces, apply flips + mirror if enabled
    faces, inferred_prefix, inferred_ext = load_cubemap_faces(
        input_path=args.input_path,
        prefix=args.prefix,
        flip_x=args.flipX,
        flip_y=args.flipY,
        mirror=mirror
    )

    # If user didn't specify a prefix, use the inferred one
    final_prefix = args.prefix if args.prefix else (inferred_prefix or "cubemap")

    # If user didn't specify output type, use inferred_ext or default to 'jpg'
    if args.type is None:
        if inferred_ext is not None and inferred_ext in ['.tga', '.png', '.jpg']:
            args.type = inferred_ext.lstrip('.')  # remove dot
        else:
            args.type = 'jpg'  # fallback

    # If user didn't specify pan => default depends on reflection map
    if args.pan is None:
        if mirror:
            args.pan = 180.0
        else:
            args.pan = 0.0

    # 2) Build the rotation matrix from pan, roll, tilt
    if abs(args.pan) > 1e-9 or abs(args.roll) > 1e-9 or abs(args.tilt) > 1e-9:
        R = euler_rotation_matrix(args.pan, args.roll, args.tilt)
    else:
        R = None

    # 3) Convert cubemap -> equirect
    equirect = cubemap_to_equirectangular(faces, out_w=out_w, out_h=out_h, R=R)
    equi_name = os.path.join(args.output_path, f"equirect_result.{args.type}")
    if equi_name.lower().endswith('.tga'):
        cv_write_tga(equi_name, equirect)
    else:
        cv2.imwrite(equi_name, equirect)
    print(f"[INFO] Wrote equirectangular to: {equi_name}")

    # 4) Make a tiny planet from that equirect
    #    We'll pick a planet size = min(out_w, out_h) typically, or some other logic.
    planet_size = min(out_w, out_h)

    planet = stereographic_tiny_planet(equirect, out_size=planet_size, fov_degrees=360)
    # Tiny planet output file: prefix + 'reflmap' + .ext
    planet_filename = f"{final_prefix}reflmap.{args.type}"
    planet_path = os.path.join(args.output_path, planet_filename)
    if planet_path.lower().endswith('.tga'):
        cv_write_tga(planet_path, planet)
    else:
        cv2.imwrite(planet_path, planet)
    print(f"[INFO] Wrote tiny planet image to: {planet_path}")

if __name__ == "__main__":
    main()