# trace generated using paraview (fixed for robust file loading, 3D camera)
import os
import glob
import sys
import subprocess
# from paraview.simple import ResetCameraClippingRange
from paraview.simple import *

#
#from paraview import servermanager as sm

#pm = sm.vtkProcessModule.GetProcessModule(
#rank = pm.GetPartitionId()
#nprocs = pm.GetNumberOfLocalPartitions()

# ---- Args ----
if len(sys.argv) < 3:
    raise SystemExit("Usage: pvbatch make_video_fixed.py  \n"
                     "Example: pvbatch --force-offscreen-rendering make_video_fixed.py "
                     "/home/heikki/0Pilvi/Lampoakku/tulokset/new_01 C")

tulos_dir = sys.argv[1]
simul_id = sys.argv[2]

if len(sys.argv) > 3:
    r_skaala = float(sys.argv[3])
else:
    r_skaala = 50.0


# Build path robustly (works with or without trailing slash)
vtk_dir = os.path.join(tulos_dir, f"vtk_{simul_id}")
print("vtk_dir:", vtk_dir, "r_skaala: ", r_skaala)

# ---- Find files ----
fluid_files = sorted(glob.glob(os.path.join(vtk_dir, "fluid_*.vtr")))
solid_files = sorted(glob.glob(os.path.join(vtk_dir, "solid_*.vtu")))

nframes = len(fluid_files)

print(f"Found {len(fluid_files)} fluid frames and {len(solid_files)} solid frames")
if not fluid_files:
    raise RuntimeError(f"No fluid_*.vtr files found in {vtk_dir}")
if not solid_files:
    raise RuntimeError(f"No solid_*.vtu files found in {vtk_dir}")

# ---- Readers ----
# Use keyword args only to avoid 'positional argument follows keyword argument'
fluid_00001vtr = XMLRectilinearGridReader(FileName=fluid_files, registrationName='fluid')
solid_00001vtu = XMLUnstructuredGridReader(FileName=solid_files, registrationName='solid')

# ---- Animation & View ----
paraview.simple._DisableFirstRenderCameraReset()
animationScene1 = GetAnimationScene()
animationScene1.UpdateAnimationUsingDataTimeSteps()
renderView1 = GetActiveViewOrCreate('RenderView')

renderView1.UseFXAA = 0          # anti-aliasing off
renderView1.Shadows = 0          # if present in your build

# ---- Show fluid ----
fluid_00001vtrDisplay = Show(fluid_00001vtr, renderView1, 'RectilinearGridRepresentation')
fluid_00001vtrDisplay.Representation = 'Surface'
fluid_00001vtrDisplay.DisableLighting = 1
renderView1.ResetCamera(True, 0.9)
renderView1.Update()

ColorBy(fluid_00001vtrDisplay, ('POINTS', 'T_fluid'))
fluid_00001vtrDisplay.RescaleTransferFunctionToDataRange(True, False)
fluid_00001vtrDisplay.SetScalarBarVisibility(renderView1, True)
t_fluidLUT = GetColorTransferFunction('T_fluid')
t_fluidPWF = GetOpacityTransferFunction('T_fluid')
t_fluidTF2D = GetTransferFunction2D('T_fluid')

# ---- Show solid ----
solid_00001vtuDisplay = Show(solid_00001vtu, renderView1, 'UnstructuredGridRepresentation')
solid_00001vtuDisplay.Representation = 'Surface'
solid_00001vtuDisplay.DisableLighting = 1
renderView1.Update()
ColorBy(solid_00001vtuDisplay, ('POINTS', 'T'))
solid_00001vtuDisplay.RescaleTransferFunctionToDataRange(True, False)
solid_00001vtuDisplay.SetScalarBarVisibility(renderView1, True)
tLUT = GetColorTransferFunction('T')
tPWF = GetOpacityTransferFunction('T')
tTF2D = GetTransferFunction2D('T')
tLUT.RescaleTransferFunction(290.0, 700.0)
tPWF.RescaleTransferFunction(290.0, 700.0)
tTF2D.RescaleTransferFunction(290.0, 700.0, 0.0, 1.0)

SetActiveSource(fluid_00001vtr)
ShowInteractiveWidgets(proxy=fluid_00001vtrDisplay)
HideInteractiveWidgets(proxy=fluid_00001vtrDisplay)
t_fluidLUT.RescaleTransferFunction(290.0, 700.0)
t_fluidPWF.RescaleTransferFunction(290.0, 700.0)
t_fluidTF2D.RescaleTransferFunction(290.0, 700.0, 0.0, 1.0)

# ---- Volume of Revolution for both ----
volumeOfRevolution1 = VolumeOfRevolution(registrationName='VolumeOfRevolution1', Input=fluid_00001vtr)
volumeOfRevolution1.Resolution = 36   # Reduce (e.g., 36) to speed up rendering
volumeOfRevolution1.SweepAngle = 360
volumeOfRevolution1Display = Show(volumeOfRevolution1, renderView1, 'UnstructuredGridRepresentation')
volumeOfRevolution1Display.Representation = 'Surface'
Hide(fluid_00001vtr, renderView1)
volumeOfRevolution1Display.SetScalarBarVisibility(renderView1, True)
renderView1.Update()

SetActiveSource(solid_00001vtu)
volumeOfRevolution2 = VolumeOfRevolution(registrationName='VolumeOfRevolution2', Input=solid_00001vtu)
volumeOfRevolution2.Resolution = 24   
volumeOfRevolution2.SweepAngle = 180
volumeOfRevolution2Display = Show(volumeOfRevolution2, renderView1, 'UnstructuredGridRepresentation')
volumeOfRevolution2Display.Representation = 'Surface'
Hide(solid_00001vtu, renderView1)
volumeOfRevolution2Display.SetScalarBarVisibility(renderView1, True)
renderView1.Update()

# Scale exaggeration
SetActiveSource(volumeOfRevolution1)
volumeOfRevolution1Display.Scale = [1.0, r_skaala, r_skaala]
volumeOfRevolution1Display.DataAxesGrid.Scale = [1.0, r_skaala, r_skaala]
volumeOfRevolution1Display.PolarAxes.Scale = [1.0, r_skaala, r_skaala]

volumeOfRevolution2Display.Scale = [1.0, r_skaala, r_skaala]
volumeOfRevolution2Display.DataAxesGrid.Scale = [1.0, r_skaala, r_skaala]
volumeOfRevolution2Display.PolarAxes.Scale = [1.0, r_skaala, r_skaala]

# Grab the scalar bars (legends) for each LUT
t_fluidBar = GetScalarBar(t_fluidLUT, renderView1)
tBar       = GetScalarBar(tLUT,        renderView1)

# Let us control size/position explicitly
for sb in (t_fluidBar, tBar):
    sb.WindowLocation     = 'Any Location'
    sb.Orientation        = 'Vertical'     # or 'Horizontal'
    sb.ScalarBarLength    = 0.25           # shorten (fraction of view height/width)
    sb.ScalarBarThickness = 8             # thinner bar
    sb.TitleFontSize      = 4             # smaller title
    sb.LabelFontSize      = 4              # smaller tick labels
    # Optional declutter:
    # sb.AddRangeLabels   = 0
    # sb.AutomaticLabelFormat = 0; sb.LabelFormat = '%.0f K'

# Position so they don't overlap (normalized [0..1] of the view)
t_fluidBar.Position = [0.85, 0.1]
tBar.Position       = [0.85, 0.6]

# renderView1.Update()
# safe_reset_clipping(renderView1)
# Fit the view to the (now scaled) data; 0.9 is a small margin
# renderView1.ResetCamera(True, 0.9)
# renderView1.ResetCameraClippingRange()

# ---- FINAL camera (3D perspective). Set this AFTER all pipeline changes. ----
renderView1.InteractionMode = '3D'
renderView1.CameraParallelProjection = 0  # perspective
renderView1.CameraPosition = [-198.0, 72.0, -117.0]
renderView1.CameraFocalPoint = [80.0, 0.0, 5.0]
renderView1.CameraViewUp = [-0.16, -0.97, -0.19]
renderView1.CameraViewAngle = 30.0
renderView1.Update()

renderView1.ResetCamera(True, 0.9)

# ---- Save animation ----
animationScene1.UpdateAnimationUsingDataTimeSteps()
nframes = len(animationScene1.TimeKeeper.TimestepValues)

out_dir = os.path.join(tulos_dir, f"frames_{simul_id}")
os.makedirs(out_dir, exist_ok=True)

for frame_num in range(nframes):
    frame_name = os.path.join(out_dir, f"frame_{frame_num:06d}.png")
    SaveScreenshot(frame_name, renderView1, ImageResolution=[1280, 720])
    animationScene1.GoToNext()

video_output = os.path.join(tulos_dir, f"video_{simul_id}.mp4")
subprocess.run([
    'ffmpeg', '-framerate', '48', 
    '-i', os.path.join(out_dir, 'frame_%06d.png'),
    '-c:v', 'libx264', '-threads', '0', '-pix_fmt', 'yuv420p',
    '-crf', '18',  # Quality: 18-23 is good (lower = better)
    video_output
])

print('finished ', video_output)
video_output = os.path.join(tulos_dir, f"video_{simul_id}.webm")
subprocess.run([
    'ffmpeg', '-framerate', '48',
    '-i', os.path.join(out_dir, 'frame_%06d.png'),
    '-c:v', 'libvpx-vp9', '-threads', '0', '-crf', '30', '-b:v', '0',
    video_output
])

print('finished ', video_output)

# out_file = os.path.join(tulos_dir, f"video_{simul_id}.ogv")
# print("Saving to:", out_file)

# SaveAnimation(
#    filename=out_file,
#    viewOrLayout=renderView1,
    # ImageResolution=[1920, 1080],  # lower to speed up (e.g., [1280, 720])
#    ImageResolution=[1280, 720],  # lower to speed up (e.g., [1280, 720])
#    FrameRate=48,
#    FrameWindow=[0, nframes-1]
# )