type color = { r : int; g : int; b : int }

type image = {
  width : int;
  height : int;
  pixels : color array array;
}

let color r g b = { r; g; b }
let bg_color = color 248 250 252
let shadow_color = color 148 163 184
let edge_color = color 203 213 225
let node_fill = color 255 255 255
let node_border = color 59 130 246
let text_color = color 30 41 59

let create_image width height =
  { width; height; pixels = Array.make_matrix height width bg_color }

let set_pixel img x y c =
  if x >= 0 && x < img.width && y >= 0 && y < img.height then
    img.pixels.(y).(x) <- c

let blend c1 c2 t =
  let t = max 0.0 (min 1.0 t) in
  let inv = 1.0 -. t in
  { r = int_of_float (float_of_int c1.r *. inv +. float_of_int c2.r *. t);
    g = int_of_float (float_of_int c1.g *. inv +. float_of_int c2.g *. t);
    b = int_of_float (float_of_int c1.b *. inv +. float_of_int c2.b *. t) }

let set_pixel_alpha img x y c alpha =
  if x >= 0 && x < img.width && y >= 0 && y < img.height then
    img.pixels.(y).(x) <- blend img.pixels.(y).(x) c alpha

let draw_shadow img cx cy radius =
  let offset = 4 in
  for dy = -radius - 8 to radius + 8 do
    for dx = -radius - 8 to radius + 8 do
      let dist = sqrt (float_of_int (dx * dx + dy * dy)) in
      if dist <= float_of_int (radius + 6) then begin
        let alpha = (1.0 -. (dist /. float_of_int (radius + 6))) *. 0.3 in
        set_pixel_alpha img (cx + dx + offset) (cy + dy + offset) shadow_color alpha
      end
    done
  done

let node_gradient_top = color 255 255 255
let node_gradient_bot = color 241 245 249

let draw_node img cx cy radius =
  for dy = -radius - 2 to radius + 2 do
    for dx = -radius - 2 to radius + 2 do
      let dist = sqrt (float_of_int (dx * dx + dy * dy)) in
      let x, y = cx + dx, cy + dy in
      let r = float_of_int radius in
      if dist <= r -. 2.5 then begin
        let grad_t = (float_of_int dy +. r) /. (2.0 *. r) in
        let fill = blend node_gradient_top node_gradient_bot grad_t in
        set_pixel img x y fill
      end
      else if dist <= r then
        set_pixel img x y node_border
      else if dist <= r +. 1.0 then
        let alpha = 1.0 -. (dist -. r) in
        set_pixel_alpha img x y node_border alpha
    done
  done

let draw_line_smooth img x0 y0 x1 y1 thickness c =
  let fx0, fy0 = float_of_int x0, float_of_int y0 in
  let fx1, fy1 = float_of_int x1, float_of_int y1 in
  let dist = sqrt ((fx1 -. fx0) ** 2.0 +. (fy1 -. fy0) ** 2.0) in
  let steps = int_of_float (dist *. 2.0) + 1 in
  for i = 0 to steps do
    let t = float_of_int i /. float_of_int steps in
    let x = fx0 +. t *. (fx1 -. fx0) in
    let y = fy0 +. t *. (fy1 -. fy0) in
    let r = float_of_int thickness in
    for dy = int_of_float (-.r -. 1.0) to int_of_float (r +. 1.0) do
      for dx = int_of_float (-.r -. 1.0) to int_of_float (r +. 1.0) do
        let d = sqrt (float_of_int (dx * dx + dy * dy)) in
        if d <= r +. 0.5 then begin
          let alpha = if d <= r -. 0.5 then 1.0 else (r +. 0.5 -. d) in
          set_pixel_alpha img (int_of_float x + dx) (int_of_float y + dy) c alpha
        end
      done
    done
  done

let draw_arc img cx cy radius start_angle end_angle thickness c =
  let steps = int_of_float (abs_float (end_angle -. start_angle) *. float_of_int radius *. 2.0) + 10 in
  for i = 0 to steps do
    let t = float_of_int i /. float_of_int steps in
    let angle = start_angle +. t *. (end_angle -. start_angle) in
    let x = float_of_int cx +. float_of_int radius *. cos angle in
    let y = float_of_int cy +. float_of_int radius *. sin angle in
    let r = float_of_int thickness in
    for dy = int_of_float (-.r -. 1.0) to int_of_float (r +. 1.0) do
      for dx = int_of_float (-.r -. 1.0) to int_of_float (r +. 1.0) do
        let d = sqrt (float_of_int (dx * dx + dy * dy)) in
        if d <= r +. 0.5 then begin
          let alpha = if d <= r -. 0.5 then 1.0 else (r +. 0.5 -. d) in
          set_pixel_alpha img (int_of_float x + dx) (int_of_float y + dy) c alpha
        end
      done
    done
  done

let pi = 3.14159265359

let draw_glyph img cx cy size ch =
  let t = max 1 (size / 8) in
  let w = size * 6 / 10 in
  let h = size in
  let x0 = cx - w / 2 in
  let x1 = cx + w / 2 in
  let y0 = cy - h / 2 in
  let y1 = cy + h / 2 in
  let ymid = cy in
  let xmid = cx in
  let r = w / 2 in
  match ch with
  | 'A' -> draw_line_smooth img x0 y1 xmid y0 t text_color;
           draw_line_smooth img xmid y0 x1 y1 t text_color;
           draw_line_smooth img (x0 + w/4) ymid (x1 - w/4) ymid t text_color
  | 'B' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 (xmid - t) y0 t text_color;
           draw_line_smooth img x0 ymid (xmid - t) ymid t text_color;
           draw_line_smooth img x0 y1 (xmid - t) y1 t text_color;
           draw_arc img (xmid - t) (y0 + (ymid - y0) / 2) ((ymid - y0) / 2) (-.pi /. 2.0) (pi /. 2.0) t text_color;
           draw_arc img (xmid - t) (ymid + (y1 - ymid) / 2) ((y1 - ymid) / 2) (-.pi /. 2.0) (pi /. 2.0) t text_color
  | 'C' -> draw_arc img xmid cy r (pi *. 0.3) (pi *. 1.7) t text_color
  | 'D' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 (xmid - r/2) y0 t text_color;
           draw_line_smooth img x0 y1 (xmid - r/2) y1 t text_color;
           draw_arc img (xmid - r/2) cy (h / 2) (-.pi /. 2.0) (pi /. 2.0) t text_color
  | 'E' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 x1 y0 t text_color;
           draw_line_smooth img x0 ymid (x1 - w/4) ymid t text_color;
           draw_line_smooth img x0 y1 x1 y1 t text_color
  | 'F' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 x1 y0 t text_color;
           draw_line_smooth img x0 ymid (x1 - w/4) ymid t text_color
  | 'G' -> draw_arc img xmid cy r (pi *. 0.3) (pi *. 1.7) t text_color;
           draw_line_smooth img xmid ymid x1 ymid t text_color;
           draw_line_smooth img x1 ymid x1 (ymid + r/2) t text_color
  | 'H' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x1 y0 x1 y1 t text_color;
           draw_line_smooth img x0 ymid x1 ymid t text_color
  | 'I' -> draw_line_smooth img xmid y0 xmid y1 t text_color;
           draw_line_smooth img (xmid - w/3) y0 (xmid + w/3) y0 t text_color;
           draw_line_smooth img (xmid - w/3) y1 (xmid + w/3) y1 t text_color
  | 'J' -> draw_line_smooth img x1 y0 x1 (y1 - r) t text_color;
           draw_arc img (x1 - r) (y1 - r) r 0.0 (pi *. 0.6) t text_color
  | 'K' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 ymid x1 y0 t text_color;
           draw_line_smooth img x0 ymid x1 y1 t text_color
  | 'L' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y1 x1 y1 t text_color
  | 'M' -> draw_line_smooth img x0 y1 x0 y0 t text_color;
           draw_line_smooth img x0 y0 xmid ymid t text_color;
           draw_line_smooth img xmid ymid x1 y0 t text_color;
           draw_line_smooth img x1 y0 x1 y1 t text_color
  | 'N' -> draw_line_smooth img x0 y1 x0 y0 t text_color;
           draw_line_smooth img x0 y0 x1 y1 t text_color;
           draw_line_smooth img x1 y1 x1 y0 t text_color
  | 'O' -> draw_arc img xmid cy r 0.0 (2.0 *. pi) t text_color
  | 'P' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 (xmid - t) y0 t text_color;
           draw_line_smooth img x0 ymid (xmid - t) ymid t text_color;
           draw_arc img (xmid - t) (y0 + (ymid - y0) / 2) ((ymid - y0) / 2) (-.pi /. 2.0) (pi /. 2.0) t text_color
  | 'Q' -> draw_arc img xmid cy r 0.0 (2.0 *. pi) t text_color;
           draw_line_smooth img xmid ymid (x1 + t) (y1 + t) t text_color
  | 'R' -> draw_line_smooth img x0 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y0 (xmid - t) y0 t text_color;
           draw_line_smooth img x0 ymid (xmid - t) ymid t text_color;
           draw_arc img (xmid - t) (y0 + (ymid - y0) / 2) ((ymid - y0) / 2) (-.pi /. 2.0) (pi /. 2.0) t text_color;
           draw_line_smooth img xmid ymid x1 y1 t text_color
  | 'S' -> draw_arc img xmid (y0 + h/4) (h/4) (pi *. 0.9) (pi *. 2.3) t text_color;
           draw_arc img xmid (y1 - h/4) (h/4) (pi *. (-0.1)) (pi *. 1.3) t text_color
  | 'T' -> draw_line_smooth img x0 y0 x1 y0 t text_color;
           draw_line_smooth img xmid y0 xmid y1 t text_color
  | 'U' -> draw_line_smooth img x0 y0 x0 (y1 - r) t text_color;
           draw_line_smooth img x1 y0 x1 (y1 - r) t text_color;
           draw_arc img xmid (y1 - r) r pi (2.0 *. pi) t text_color
  | 'V' -> draw_line_smooth img x0 y0 xmid y1 t text_color;
           draw_line_smooth img xmid y1 x1 y0 t text_color
  | 'W' -> draw_line_smooth img x0 y0 (x0 + w/4) y1 t text_color;
           draw_line_smooth img (x0 + w/4) y1 xmid (ymid + h/6) t text_color;
           draw_line_smooth img xmid (ymid + h/6) (x1 - w/4) y1 t text_color;
           draw_line_smooth img (x1 - w/4) y1 x1 y0 t text_color
  | 'X' -> draw_line_smooth img x0 y0 x1 y1 t text_color;
           draw_line_smooth img x0 y1 x1 y0 t text_color
  | 'Y' -> draw_line_smooth img x0 y0 xmid ymid t text_color;
           draw_line_smooth img x1 y0 xmid ymid t text_color;
           draw_line_smooth img xmid ymid xmid y1 t text_color
  | 'Z' -> draw_line_smooth img x0 y0 x1 y0 t text_color;
           draw_line_smooth img x1 y0 x0 y1 t text_color;
           draw_line_smooth img x0 y1 x1 y1 t text_color
  | '0' -> draw_arc img xmid cy r 0.0 (2.0 *. pi) t text_color;
           draw_line_smooth img (x0 + w/4) (y1 - h/4) (x1 - w/4) (y0 + h/4) t text_color
  | '1' -> draw_line_smooth img xmid y0 xmid y1 t text_color;
           draw_line_smooth img (xmid - w/3) (y0 + h/4) xmid y0 t text_color;
           draw_line_smooth img (xmid - w/3) y1 (xmid + w/3) y1 t text_color
  | '2' -> draw_arc img xmid (y0 + h/3) (h/3) (pi *. 1.2) (pi *. 2.1) t text_color;
           draw_line_smooth img (x1 - t/2) (y0 + h/3) x0 y1 t text_color;
           draw_line_smooth img x0 y1 x1 y1 t text_color
  | '3' -> draw_arc img xmid (y0 + h/4) (h/4) (pi *. 1.1) (pi *. 2.0) t text_color;
           draw_arc img xmid (y1 - h/4) (h/4) 0.0 (pi *. 0.9) t text_color
  | '4' -> draw_line_smooth img (x1 - w/4) y0 (x1 - w/4) y1 t text_color;
           draw_line_smooth img (x1 - w/4) y0 x0 (ymid + h/8) t text_color;
           draw_line_smooth img x0 (ymid + h/8) x1 (ymid + h/8) t text_color
  | '5' -> draw_line_smooth img x1 y0 x0 y0 t text_color;
           draw_line_smooth img x0 y0 x0 ymid t text_color;
           draw_line_smooth img x0 ymid (xmid - t) ymid t text_color;
           draw_arc img (xmid - t) (ymid + (y1 - ymid) / 2) ((y1 - ymid) / 2) (-.pi /. 2.0) (pi *. 0.7) t text_color
  | '6' -> draw_arc img xmid (y1 - h/3) (h/3) 0.0 (2.0 *. pi) t text_color;
           draw_arc img (x1 + w/5) (y0 + h/3) (h/2) (pi *. 0.7) (pi *. 1.0) t text_color
  | '7' -> draw_line_smooth img x0 y0 x1 y0 t text_color;
           draw_line_smooth img x1 y0 (x0 + w/4) y1 t text_color
  | '8' -> draw_arc img xmid (y0 + h/4) (h/4 - t/2) 0.0 (2.0 *. pi) t text_color;
           draw_arc img xmid (y1 - h/4) (h/4) 0.0 (2.0 *. pi) t text_color
  | '9' -> draw_arc img xmid (y0 + h/3) (h/3) 0.0 (2.0 *. pi) t text_color;
           draw_arc img (x0 - w/5) (y1 - h/3) (h/2) (-.pi *. 0.3) 0.0 t text_color
  | '-' -> draw_line_smooth img (x0 + w/4) ymid (x1 - w/4) ymid t text_color
  | _ -> ()

let draw_text img cx cy radius text =
  let len = String.length text in
  let max_w = float_of_int radius *. 1.3 in
  let max_h = float_of_int radius *. 0.7 in
  let char_size_w = int_of_float (max_w /. (float_of_int len *. 0.7)) in
  let char_size_h = int_of_float max_h in
  let char_size = min char_size_w char_size_h in
  let char_w = char_size * 7 / 10 in
  let total_w = len * char_w in
  let start_x = cx - total_w / 2 + char_w / 2 in
  String.iteri (fun i ch ->
    let upper = if ch >= 'a' && ch <= 'z' then Char.chr (Char.code ch - 32) else ch in
    draw_glyph img (start_x + i * char_w) cy char_size upper
  ) text

let bezier_point p0 p1 p2 p3 t =
  let inv = 1.0 -. t in
  let inv2 = inv *. inv in
  let inv3 = inv2 *. inv in
  let t2 = t *. t in
  let t3 = t2 *. t in
  inv3 *. p0 +. 3.0 *. inv2 *. t *. p1 +. 3.0 *. inv *. t2 *. p2 +. t3 *. p3

let draw_bezier_aa img x0 y0 x1 y1 x2 y2 x3 y3 thickness c =
  let steps = int_of_float (sqrt ((x3-.x0)**2.0 +. (y3-.y0)**2.0) *. 3.0) in
  let steps = max steps 80 in
  for i = 0 to steps do
    let t = float_of_int i /. float_of_int steps in
    let x = bezier_point x0 x1 x2 x3 t in
    let y = bezier_point y0 y1 y2 y3 t in
    for dx = -thickness - 1 to thickness + 1 do
      for dy = -thickness - 1 to thickness + 1 do
        let dist = sqrt (float_of_int (dx * dx + dy * dy)) in
        if dist <= float_of_int thickness +. 0.8 then begin
          let alpha = if dist <= float_of_int thickness then 1.0
                      else 1.0 -. (dist -. float_of_int thickness) /. 0.8 in
          set_pixel_alpha img (int_of_float x + dx) (int_of_float y + dy) c alpha
        end
      done
    done
  done

let draw_curved_edge img x0 y0 x1 y1 =
  let fx0, fy0 = float_of_int x0, float_of_int y0 in
  let fx1, fy1 = float_of_int x1, float_of_int y1 in
  let ctrl_y = fy0 +. (fy1 -. fy0) *. 0.4 in
  draw_bezier_aa img fx0 fy0 fx0 ctrl_y fx1 (fy1 -. (fy1 -. fy0) *. 0.4) fx1 fy1 2 edge_color

let write_ppm filename img =
  let oc = open_out filename in
  Printf.fprintf oc "P3\n%d %d\n255\n" img.width img.height;
  for y = 0 to img.height - 1 do
    for x = 0 to img.width - 1 do
      let p = img.pixels.(y).(x) in
      Printf.fprintf oc "%d %d %d " p.r p.g p.b
    done;
    output_char oc '\n'
  done;
  close_out oc

let crc32_table =
  Array.init 256 (fun n ->
    let c = ref (Int32.of_int n) in
    for _ = 0 to 7 do
      if Int32.(logand !c 1l = 1l) then
        c := Int32.(logxor (shift_right_logical !c 1) 0xEDB88320l)
      else
        c := Int32.shift_right_logical !c 1
    done;
    !c)

let crc32 data =
  let crc = ref 0xFFFFFFFFl in
  Bytes.iter (fun b ->
    let idx = Int32.(to_int (logand (logxor !crc (of_int (Char.code b))) 0xFFl)) in
    crc := Int32.(logxor (shift_right_logical !crc 8) crc32_table.(idx))
  ) data;
  Int32.logxor !crc 0xFFFFFFFFl

let adler32 data =
  let a = ref 1 and b = ref 0 in
  Bytes.iter (fun c ->
    a := (!a + Char.code c) mod 65521;
    b := (!b + !a) mod 65521
  ) data;
  Int32.(logor (shift_left (of_int !b) 16) (of_int !a))

let write_png filename img =
  let oc = open_out_bin filename in
  let write_byte b = output_byte oc b in
  let write_int32_be n =
    write_byte (Int32.to_int (Int32.shift_right_logical n 24) land 0xFF);
    write_byte (Int32.to_int (Int32.shift_right_logical n 16) land 0xFF);
    write_byte (Int32.to_int (Int32.shift_right_logical n 8) land 0xFF);
    write_byte (Int32.to_int n land 0xFF)
  in
  let write_chunk chunk_type data =
    write_int32_be (Int32.of_int (Bytes.length data));
    let type_and_data = Bytes.cat (Bytes.of_string chunk_type) data in
    output_bytes oc type_and_data;
    write_int32_be (crc32 type_and_data)
  in
  output_string oc "\137PNG\r\n\026\n";
  let ihdr = Bytes.create 13 in
  Bytes.set_int32_be ihdr 0 (Int32.of_int img.width);
  Bytes.set_int32_be ihdr 4 (Int32.of_int img.height);
  Bytes.set ihdr 8 '\008'; Bytes.set ihdr 9 '\002';
  Bytes.set ihdr 10 '\000'; Bytes.set ihdr 11 '\000'; Bytes.set ihdr 12 '\000';
  write_chunk "IHDR" ihdr;
  let raw_size = img.height * (1 + img.width * 3) in
  let raw = Bytes.create raw_size in
  let pos = ref 0 in
  for y = 0 to img.height - 1 do
    Bytes.set raw !pos '\000'; incr pos;
    for x = 0 to img.width - 1 do
      let p = img.pixels.(y).(x) in
      Bytes.set raw !pos (Char.chr p.r); incr pos;
      Bytes.set raw !pos (Char.chr p.g); incr pos;
      Bytes.set raw !pos (Char.chr p.b); incr pos
    done
  done;
  let max_block = 65535 in
  let num_blocks = (raw_size + max_block - 1) / max_block in
  let zlib_size = 2 + raw_size + num_blocks * 5 + 4 in
  let zlib = Bytes.create zlib_size in
  Bytes.set zlib 0 '\x78'; Bytes.set zlib 1 '\x01';
  let zpos = ref 2 in
  let rpos = ref 0 in
  for i = 0 to num_blocks - 1 do
    let remaining = raw_size - !rpos in
    let block_len = min max_block remaining in
    let is_last = i = num_blocks - 1 in
    Bytes.set zlib !zpos (if is_last then '\x01' else '\x00'); incr zpos;
    Bytes.set zlib !zpos (Char.chr (block_len land 0xFF)); incr zpos;
    Bytes.set zlib !zpos (Char.chr ((block_len lsr 8) land 0xFF)); incr zpos;
    let nlen = block_len lxor 0xFFFF in
    Bytes.set zlib !zpos (Char.chr (nlen land 0xFF)); incr zpos;
    Bytes.set zlib !zpos (Char.chr ((nlen lsr 8) land 0xFF)); incr zpos;
    Bytes.blit raw !rpos zlib !zpos block_len;
    zpos := !zpos + block_len;
    rpos := !rpos + block_len
  done;
  let adler = adler32 raw in
  Bytes.set_int32_be zlib !zpos adler;
  write_chunk "IDAT" zlib;
  write_chunk "IEND" Bytes.empty;
  close_out oc

let render_tree ~filename ~node_radius ~h_scale ~v_scale ~padding nodes edges =
  let min_x = List.fold_left (fun acc (_, x, _) -> min acc x) 0.0 nodes in
  let max_x = List.fold_left (fun acc (_, x, _) -> max acc x) 0.0 nodes in
  let max_depth = List.fold_left (fun acc (_, _, d) -> max acc d) 0 nodes in

  let width = int_of_float ((max_x -. min_x) *. h_scale) + 2 * padding + 2 * node_radius in
  let height = (max_depth + 1) * v_scale + 2 * padding in

  let img = create_image width height in

  let to_pixel_x x = int_of_float ((x -. min_x) *. h_scale) + padding + node_radius in
  let to_pixel_y depth = depth * v_scale + padding + node_radius in

  List.iter (fun (px, py, cx, cy) ->
    let px', py' = to_pixel_x px, to_pixel_y py in
    let cx', cy' = to_pixel_x cx, to_pixel_y cy in
    draw_curved_edge img px' py' cx' cy'
  ) edges;

  List.iter (fun (_, x, depth) ->
    let px, py = to_pixel_x x, to_pixel_y depth in
    draw_shadow img px py node_radius
  ) nodes;

  List.iter (fun (_, x, depth) ->
    let px, py = to_pixel_x x, to_pixel_y depth in
    draw_node img px py node_radius
  ) nodes;

  let is_png = String.length filename >= 4 &&
    String.sub filename (String.length filename - 4) 4 = ".png" in
  if is_png then write_png filename img else write_ppm filename img;
  Printf.printf "Wrote %s (%dx%d)\n" filename width height