(* ================================================================
   Réseau de Neurones Convolutionnel (CNN)
   
   Un CNN apprend à classifier des motifs sur une grille 2D.
   Il combine convolution, activation ReLU, pooling et couche
   entièrement connectée.
   
   Architecture :
     Entrée (6×6)
       → Conv2D (2 filtres 3×3) → ReLU → MaxPool (2×2)
       → Aplatir → FullyConnected (2 classes)
   
   Exemple : distinguer bandes verticales / horizontales
   ================================================================ *)

#directory "/home/thiry/Bureau/ML"
#load "graphiques.cmo"
open Graphiques

let _ = Random.self_init ()

(* ----------------------------------------------------------------- *)
(*  1. Primitives mathématiques                                       *)
(* ----------------------------------------------------------------- *)

let dot xs ys = List.fold_left2 (fun s x y -> s +. x *. y) 0. xs ys
let vec_add = List.map2 (+.)
let vec_sub = List.map2 (-.)
let vec_smul a = List.map (fun x -> a *. x)
let vec_mul = List.map2 ( *. )
let mat_mul m v = List.map (fun row -> dot row v) m

let relu x  = if x > 0. then x else 0.
let drelu x = if x > 0. then 1. else 0.

(* ----------------------------------------------------------------- *)
(*  2. Types de couches                                               *)
(* ----------------------------------------------------------------- *)

(* Couche de convolution 2D (entrée monochrome) *)
type conv2d = {
  mutable noyaux : float list list list;  (* n_filtres × hk × wk *)
  mutable biais  : float list;            (* n_filtres *)
  n_filtres : int;
  hk : int; wk : int;
}

(* Couche de pooling : max sur fenêtre 2×2 stride 2 *)
type pool = Pool

(* Couche fully connected *)
type fc = {
  mutable poids : float list list;  (* n_sortie × n_entree *)
  mutable biais : float list;       (* n_sortie *)
}

(* ----------------------------------------------------------------- *)
(*  3. Initialisation                                                 *)
(* ----------------------------------------------------------------- *)

let rand () = Random.float 2. -. 1.

let cree_conv2d n_filtres hk wk =
  { noyaux = List.init n_filtres (fun _ ->
      List.init hk (fun _ -> List.init wk (fun _ -> rand () /. sqrt (float (hk * wk)))));
    biais  = List.init n_filtres (fun _ -> 0.);
    n_filtres; hk; wk }

let cree_fc n_entree n_sortie =
  { poids = List.init n_sortie (fun _ ->
      List.init n_entree (fun _ -> rand () /. sqrt (float n_entree)));
    biais = List.init n_sortie (fun _ -> 0.) }

(* ----------------------------------------------------------------- *)
(*  4. Convolution 2D (pas avant)                                     *)
(* ----------------------------------------------------------------- *)

let conv2d_forward (c : conv2d) image =
  (* image : hi × wi  (liste de listes) *)
  let hi = List.length image in
  let wi = List.length (List.hd image) in
  let ho = hi - c.hk + 1 in
  let wo = wi - c.wk + 1 in
  List.init c.n_filtres (fun k ->
    let noyau = List.nth c.noyaux k in
    List.init ho (fun i ->
      List.init wo (fun j ->
        let somme = ref (List.nth c.biais k) in
        for di = 0 to c.hk - 1 do
          for dj = 0 to c.wk - 1 do
            let vi = List.nth (List.nth image (i + di)) (j + dj) in
            let wk_val = List.nth (List.nth noyau di) dj in
            somme := !somme +. vi *. wk_val
          done
        done;
        !somme
      )
    )
  )
  (* résultat : n_filtres × ho × wo *)

(* ----------------------------------------------------------------- *)
(*  5. Max Pooling 2×2 stride 2 (pas avant)                           *)
(* ----------------------------------------------------------------- *)

(* Version avec cartes -> pour chaque filtre on pool séparément *)
let max_pool_forward cartes =
  (* cartes : n_filtres × hi × wi *)
  let n_filtres = List.length cartes in
  let hi = List.length (List.hd cartes) in
  let wi = List.length (List.hd (List.hd cartes)) in
  let ho = hi / 2 in
  let wo = wi / 2 in
  List.map (fun carte ->
    List.init ho (fun i ->
      List.init wo (fun j ->
        let a = List.nth (List.nth carte (2*i)) (2*j) in
        let b = List.nth (List.nth carte (2*i)) (2*j+1) in
        let c = List.nth (List.nth carte (2*i+1)) (2*j) in
        let d = List.nth (List.nth carte (2*i+1)) (2*j+1) in
        max (max a b) (max c d)
      )
    )
  ) cartes

(* Retourne aussi les indices du max pour la rétroprop *)
let max_pool_forward_idx cartes =
  let n_filtres = List.length cartes in
  let hi = List.length (List.hd cartes) in
  let wi = List.length (List.hd (List.hd cartes)) in
  let ho = hi / 2 in
  let wo = wi / 2 in
  List.map (fun carte ->
    List.init ho (fun i ->
      List.init wo (fun j ->
        let positions = [
          (0,0, List.nth (List.nth carte (2*i)) (2*j));
          (0,1, List.nth (List.nth carte (2*i)) (2*j+1));
          (1,0, List.nth (List.nth carte (2*i+1)) (2*j));
          (1,1, List.nth (List.nth carte (2*i+1)) (2*j+1));
        ] in
        let (di, dj, _) = List.fold_left (fun (ai,aj,av) (bi,bj,bv) ->
          if bv > av then (bi,bj,bv) else (ai,aj,av)
        ) (0,0, neg_infinity) positions in
        (di, dj)
      )
    )
  ) cartes

(* ----------------------------------------------------------------- *)
(*  6. Softmax et cross-entropy                                       *)
(* ----------------------------------------------------------------- *)

let softmax xs =
  let m = List.fold_left max 0. xs in
  let exps = List.map (fun x -> exp (x -. m)) xs in
  let sum = List.fold_left (+.) 0. exps in
  List.map (fun e -> e /. sum) exps

let cross_entropy pred cible =
  List.fold_left2 (fun s p c ->
    s -. (if c = 0 then 0. else log p)
  ) 0. pred cible

(* ----------------------------------------------------------------- *)
(*  7. Passage avant complet                                          *)
(* ----------------------------------------------------------------- *)

type cache = {
  image   : float list list;
  conv_z  : float list list list; (* avant ReLU *)
  conv_a  : float list list list; (* après ReLU *)
  pool_idx: (int * int) list list list;
  pool_out: float list list list;
  plat    : float list;
  logits  : float list;
  probs   : float list;
}

let forward (c : conv2d) (f : fc) image =
  (* 1. Convolution *)
  let conv_z = conv2d_forward c image in
  (* 2. ReLU *)
  let conv_a = List.map (fun carte ->
    List.map (fun ligne -> List.map relu ligne) carte
  ) conv_z in
  (* 3. Max Pooling 2×2 *)
  let (pool_idx, pool_out) =
    let idx = max_pool_forward_idx conv_a in
    (idx, max_pool_forward conv_a) in
  (* 4. Aplatir *)
  let plat = List.flatten (List.flatten pool_out) in
  (* 5. Fully Connected *)
  let logits = vec_add (mat_mul f.poids plat) f.biais in
  let probs = softmax logits in
  { image; conv_z; conv_a; pool_idx; pool_out; plat; logits; probs }

(* ----------------------------------------------------------------- *)
(*  8. Rétropropagation                                               *)
(* ----------------------------------------------------------------- *)

(* Gradient pour la couche fully connected *)
let backward_fc (f : fc) (c : cache) cible taux =
  (* dL/dlogits = probs - one_hot(cible) *)
  let n_c = List.length c.probs in
  let d_logits = List.init n_c (fun i ->
    List.nth c.probs i -. if i = cible then 1. else 0.
  ) in
  let d_poids = List.map (fun d_j ->
    List.map (fun x_i -> d_j *. x_i) c.plat
  ) d_logits in
  let d_biais = d_logits in
  let d_plat = List.init (List.length c.plat) (fun i ->
    List.fold_left2 (fun s w_ji d_j -> s +. w_ji *. d_j) 0.
      (List.map (fun row -> List.nth row i) f.poids) d_logits
  ) in
  (* Mise à jour des poids FC *)
  f.poids <- List.map2 (fun row d_row ->
    vec_sub row (vec_smul taux d_row)
  ) f.poids d_poids;
  f.biais <- vec_sub f.biais (vec_smul taux d_biais);
  d_plat

(* Gradient pour le max pooling (2×2) *)
let backward_pool (c : cache) d_plat =
  let n_filtres = List.length c.conv_a in
  let hi = List.length (List.hd c.conv_a) in
  let wi = List.length (List.hd (List.hd c.conv_a)) in
  let d_carte = Array.init n_filtres (fun _ ->
    Array.init hi (fun _ -> Array.make wi 0.)
  ) in
  let idx_plateau = ref 0 in
  List.iteri (fun k idxs ->
    List.iteri (fun i ligne ->
      List.iteri (fun j (di, dj) ->
        let val_d = List.nth d_plat !idx_plateau in
        d_carte.(k).(2*i+di).(2*j+dj) <-
          d_carte.(k).(2*i+di).(2*j+dj) +. val_d;
        incr idx_plateau
      ) ligne
    ) idxs
  ) c.pool_idx;
  List.init n_filtres (fun k ->
    List.init hi (fun i ->
      List.init wi (fun j -> d_carte.(k).(i).(j))
    )
  )

(* Gradient pour ReLU *)
let backward_relu (c : cache) d_carte =
  List.map2 (fun conv_z_carte d_carte_carte ->
    List.map2 (fun ligne_z ligne_d ->
      List.map2 (fun z d -> if z > 0. then d else 0.) ligne_z ligne_d
    ) conv_z_carte d_carte_carte
  ) c.conv_z d_carte

(* Gradient pour la convolution *)
let backward_conv (c : conv2d) (cache : cache) d_conv taux =
  let hi = List.length cache.image in
  let wi = List.length (List.hd cache.image) in
  (* dL/dnoyau *)
  let d_noyaux = List.init c.n_filtres (fun k ->
    List.init c.hk (fun di ->
      List.init c.wk (fun dj ->
        let somme = ref 0. in
        for i = 0 to hi - c.hk do
          for j = 0 to wi - c.wk do
            let dy = List.nth (List.nth (List.nth d_conv k) i) j in
            let x  = List.nth (List.nth cache.image (i + di)) (j + dj) in
            somme := !somme +. dy *. x
          done
        done;
        !somme
      )
    )
  ) in
  let d_biais = List.map (fun carte ->
    List.fold_left (fun s ligne ->
      s +. List.fold_left (+.) 0. ligne
    ) 0. carte
  ) d_conv in
  (* dL/dimage (pour info, pas utilisée ici) *)
  let d_image = List.init hi (fun i ->
    List.init wi (fun j ->
      let somme = ref 0. in
      for k = 0 to c.n_filtres - 1 do
        let noyau = List.nth c.noyaux k in
        for di = 0 to c.hk - 1 do
          for dj = 0 to c.wk - 1 do
            let pi = i - di in
            let pj = j - dj in
            if pi >= 0 && pi < hi - c.hk + 1 &&
               pj >= 0 && pj < wi - c.wk + 1 then
              let dy = List.nth (List.nth (List.nth d_conv k) pi) pj in
              let w  = List.nth (List.nth noyau di) dj in
              somme := !somme +. dy *. w
          done
        done
      done;
      !somme
    )
  ) in
  (* Mise à jour des noyaux *)
  c.noyaux <- List.map2 (fun noyau d_no ->
    List.map2 (fun ligne d_ligne ->
      vec_sub ligne (vec_smul taux d_ligne)
    ) noyau d_no
  ) c.noyaux d_noyaux;
  c.biais <- vec_sub c.biais (vec_smul taux d_biais);
  d_image

(* Passage complet backward + mise à jour *)
let backward (c : conv2d) (f : fc) (cache : cache) cible taux =
  let d_plat = backward_fc f cache cible taux in
  let d_pool = backward_pool cache d_plat in
  let d_relu = backward_relu cache d_pool in
  ignore (backward_conv c cache d_relu taux)

(* ----------------------------------------------------------------- *)
(*  9. Création des jeux de données                                   *)
(* ----------------------------------------------------------------- *)

(* Génère une image avec une bande verticale (classe 0) *)
let bande_verticale () =
  let decal = Random.int 3 in (* 0 .. 2 *)
  List.init 6 (fun i ->
    List.init 6 (fun j ->
      if j >= decal && j < decal + 2 then 1. else 0.
    )
  )

(* Génère une image avec une bande horizontale (classe 1) *)
let bande_horizontale () =
  let decal = Random.int 3 in
  List.init 6 (fun i ->
    List.init 6 (fun j ->
      if i >= decal && i < decal + 2 then 1. else 0.
    )
  )

let creer_donnees n_par_classe =
  let donnees = ref [] in
  for _ = 1 to n_par_classe do
    donnees := (bande_verticale (), 0) :: !donnees;
    donnees := (bande_horizontale (), 1) :: !donnees
  done;
  !donnees

(* ----------------------------------------------------------------- *)
(* 10. Entraînement                                                   *)
(* ----------------------------------------------------------------- *)

let entrainer (c : conv2d) (f : fc) donnees taux n_iter =
  let n = float_of_int (List.length donnees) in
  for iter = 1 to n_iter do
    let bon = ref 0 in
    let cout = ref 0. in
    List.iter (fun (image, cible) ->
      let cache = forward c f image in
      let pred = List.hd (List.sort (fun a b -> compare (snd b) (snd a))
        (List.mapi (fun i p -> (i, p)) cache.probs)) in
      let (pred_classe, _) = pred in
      if pred_classe = cible then incr bon;
      cout := !cout +. cross_entropy cache.probs
        (List.init (List.length cache.probs) (fun i -> if i = cible then 1 else 0));
      backward c f cache cible (taux /. n)
    ) donnees;
    if iter mod 20 = 0 || iter = 1 then
      Printf.printf "  iter %4d : precision = %d/%d (%.1f%%)  cout = %.4f\n%!"
        iter !bon (List.length donnees)
        (100. *. float !bon /. n) (!cout /. n)
  done

(* ----------------------------------------------------------------- *)
(* 11. Visualisation d'une image                                      *)
(* ----------------------------------------------------------------- *)

let affiche_image img =
  List.iter (fun ligne ->
    Printf.printf "    ";
    List.iter (fun p ->
      if p > 0.5 then Printf.printf "██" else Printf.printf "  "
    ) ligne;
    Printf.printf "\n%!"
  ) img

(* ----------------------------------------------------------------- *)
(* 12. Démo                                                           *)
(* ----------------------------------------------------------------- *)

let () =
  Printf.printf "=== CNN : Classification de motifs ===\n\n%!";
  Printf.printf "  Architecture :\n";
  Printf.printf "    Entrée : 6×6\n";
  Printf.printf "    Conv2D : 2 filtres 3×3 + ReLU\n";
  Printf.printf "    MaxPool: 2×2 (stride 2)\n";
  Printf.printf "    FC     : 8 → 2 (softmax)\n\n%!";

  let conv = cree_conv2d 2 3 3 in
  let fc   = cree_fc 8 2 in
  let donnees = creer_donnees 30 in

  Printf.printf "  %d images d'entraînement\n%!" (List.length donnees);

  Printf.printf "\n  Exemples de motifs :\n%!";
  Printf.printf "\n  --- Bande verticale (classe 0) ---\n%!";
  affiche_image (List.hd (List.filter (fun (_, c) -> c = 0) donnees) |> fst);
  Printf.printf "\n  --- Bande horizontale (classe 1) ---\n%!";
  affiche_image (List.hd (List.filter (fun (_, c) -> c = 1) donnees) |> fst);

  Printf.printf "\n  Entraînement en cours...\n%!";
  entrainer conv fc donnees 1.0 100;

  Printf.printf "\n  Test sur nouveaux motifs :\n%!";
  let test_v = bande_verticale () in
  let test_h = bande_horizontale () in
  let cache_v = forward conv fc test_v in
  let cache_h = forward conv fc test_h in
  Printf.printf "\n  Bande verticale → probs = [%.3f ; %.3f] (classe %d)\n%!"
    (List.nth cache_v.probs 0) (List.nth cache_v.probs 1)
    (fst (List.hd (List.sort (fun a b -> compare (snd b) (snd a))
      (List.mapi (fun i p -> (i, p)) cache_v.probs))));
  Printf.printf "  Bande horizontale → probs = [%.3f ; %.3f] (classe %d)\n%!"
    (List.nth cache_h.probs 0) (List.nth cache_h.probs 1)
    (fst (List.hd (List.sort (fun a b -> compare (snd b) (snd a))
      (List.mapi (fun i p -> (i, p)) cache_h.probs))));

  Printf.printf "\n  Cartes de convolution apprises :\n%!";
  List.iteri (fun k noyau ->
    Printf.printf "  Filtre %d :\n%!" k;
    List.iter (fun ligne ->
      Printf.printf "    ";
      List.iter (fun v -> Printf.printf "%+.2f " v) ligne;
      Printf.printf "\n%!"
    ) noyau
  ) conv.noyaux

(* ----------------------------------------------------------------- *)
(* 13. Graphiques SVG                                                 *)
(* ----------------------------------------------------------------- *)

let () =
  (* Re-entrainement avec historique *)
  let conv = cree_conv2d 2 3 3 in
  let fc   = cree_fc 8 2 in
  let donnees = creer_donnees 30 in
  let n = float_of_int (List.length donnees) in
  let historique = ref [] in
  for iter = 1 to 100 do
    let bon = ref 0 in
    let cout = ref 0. in
    List.iter (fun (image, cible) ->
      let cache = forward conv fc image in
      let pred = fst (List.hd (List.sort (fun a b -> compare (snd b) (snd a))
        (List.mapi (fun i p -> (i, p)) cache.probs))) in
      if pred = cible then incr bon;
      cout := !cout +. cross_entropy cache.probs
        (List.init (List.length cache.probs) (fun i -> if i = cible then 1 else 0));
      backward conv fc cache cible (1.0 /. n)
    ) donnees;
    if iter mod 5 = 0 then
      historique := (float_of_int iter, 100. *. float !bon /. n) :: !historique
  done;

  courbe_entrainement ~fichier:"cnn_precision.svg"
    ~titre:"Entrainement du CNN"
    ~xlab:"Iteration" ~ylab:"Precision (%)"
    [{ points = List.rev !historique; legende = "Precision" }];

  (* Carte de chaleur des filtres appris *)
  let () =
    let () = Printf.printf "  -> Sauvegarde des filtres...\n%!" in
    List.iteri (fun k noyau ->
      carte_chaleur ~fichier:(Printf.sprintf "cnn_filtre_%d.svg" k)
        ~titre:(Printf.sprintf "Filtre %d" k) noyau
    ) conv.noyaux
  in ()