/-
  Convolution.lean — Cross-theorem pattern composition via comonadic convolution.

  Takes extractable clusters from different theorems, finds compatible slots
  (same type, both varying), and generates a composed pattern that unifies
  the shared structure.

  ══════════════════════════════════════════════════════════
  COMMANDS
  ══════════════════════════════════════════════════════════

    #patternCompose thm1 thm2
    #patternPreview thm1 c1 thm2 c2
    #patternExecute thm1 c1 thm2 c2 as "composedName"
-/

import Infoductor.Comonad.ComonadFinder
import Infoductor.Comonad.ExtractDefn

open Lean Meta Elab Command ExtractDefn

namespace Convolution

-- ════════════════════════════════════════════════════════════════════════════
-- § 1  Core data structures
-- ════════════════════════════════════════════════════════════════════════════

/-- Alignment between two slot positions from different clusters. -/
structure SlotAlignment where
  pos1       : Nat
  pos2       : Nat
  typeStr1   : String
  typeStr2   : String
  unifiable  : Bool
  deriving Repr, BEq, Hashable

/-- A candidate for convolution: two extractable clusters that may compose. -/
structure ConvolutionCandidate where
  theorem1   : String
  theorem2   : String
  cluster1Id : Nat
  cluster2Id : Nat
  shape1     : String
  shape2     : String
  alignments : Array SlotAlignment
  score      : Float
  deriving Repr

/-- Result of composing two patterns. -/
structure ComposedPattern where
  name         : String
  source1      : String × Nat
  source2      : String × Nat
  alignments   : Array SlotAlignment
  combinedShape: String
  paramCount   : Nat
  leanCommand  : String
  newGraphId   : ContentHash
  deriving Repr

-- ════════════════════════════════════════════════════════════════════════════
-- § 2  Type compatibility checking
-- ════════════════════════════════════════════════════════════════════════════

/-- Check if two type strings are structurally similar. -/
private def typesCompatible (t1 t2 : String) : Bool :=
  -- Lean 4 v4.30 dropped `String.containsSubstr`; use a small inline
  -- substring check instead.
  let containsArrow (s : String) : Bool := (s.splitOn "→").length > 1
  if t1 == t2 then true
  else if containsArrow t1 && containsArrow t2 then
    let ret1 := t1.splitOn "→" |>.getLast?
    let ret2 := t2.splitOn "→" |>.getLast?
    ret1 == ret2
  else if t1 == "Prop" && t2 == "Prop" then true
  else
    let base1 := t1.splitOn " " |>.head?
    let base2 := t2.splitOn " " |>.head?
    base1.isSome && base1 == base2

/-- Compute a compatibility score based on alignments. -/
private def computeScore (alignments : Array SlotAlignment) (totalSlots : Nat) : Float :=
  if totalSlots == 0 then 0.0
  else
    let unifiableCount := (alignments.filter (·.unifiable)).size
    unifiableCount.toFloat / totalSlots.toFloat

-- ════════════════════════════════════════════════════════════════════════════
-- § 3  Find compatible slots between two clusters
-- ════════════════════════════════════════════════════════════════════════════

/-- Build type map for varying slots from a proposal. -/
private def getVarSlotTypes (proposal : ExtractionProposal)
    (richNodes : Array RichNode) : Std.HashMap Nat String :=
  let nodeMap : Std.HashMap ContentHash String :=
    richNodes.foldl (fun m n => m.insert n.contentId n.typeStr) {}
  proposal.varSlots.foldl (fun m slot =>
    let typeStr := slot.varIds[0]?.bind (fun id => nodeMap.get? id) |>.getD "_"
    m.insert slot.pos typeStr
  ) {}

/-- Find alignable slots between two extraction proposals. -/
def findCompatibleSlots
    (prop1 : ExtractionProposal) (types1 : Std.HashMap Nat String)
    (prop2 : ExtractionProposal) (types2 : Std.HashMap Nat String)
    : Array SlotAlignment := Id.run do
  let mut alignments : Array SlotAlignment := #[]
  for vs1 in prop1.varSlots do
    for vs2 in prop2.varSlots do
      let t1 := types1.getD vs1.pos "_"
      let t2 := types2.getD vs2.pos "_"
      let unifiable := typesCompatible t1 t2
      alignments := alignments.push {
        pos1      := vs1.pos
        pos2      := vs2.pos
        typeStr1  := t1
        typeStr2  := t2
        unifiable := unifiable
      }
  return alignments

-- ════════════════════════════════════════════════════════════════════════════
-- § 4  Find convolution candidates across theorems
-- ════════════════════════════════════════════════════════════════════════════

/-- Find all convoluble cluster pairs between two theorems. -/
def findCandidates (thm1 thm2 : Name) : MetaM (Array ConvolutionCandidate) := do
  let graph1 ← findComonadsCore thm1
  let graph2 ← findComonadsCore thm2

  let ci1 ← getConstInfo thm1
  let ci2 ← getConstInfo thm2
  let val1 ← match ci1.value? with | some e => pure e | none => throwError "no value for {thm1}"
  let val2 ← match ci2.value? with | some e => pure e | none => throwError "no value for {thm2}"
  let entries1 ← Mathlib.Explode.explode val1
  let entries2 ← Mathlib.Explode.explode val2
  let richNodes1 ← buildRichIR entries1
  let richNodes2 ← buildRichIR entries2

  let extractable1 := graph1.clusters.filter (·.extractable)
  let extractable2 := graph2.clusters.filter (·.extractable)

  let mut candidates : Array ConvolutionCandidate := #[]

  for c1 in extractable1 do
    for c2 in extractable2 do
      let prop1 := analyzeCluster c1
      let prop2 := analyzeCluster c2
      let types1 := getVarSlotTypes prop1 richNodes1
      let types2 := getVarSlotTypes prop2 richNodes2
      let alignments := findCompatibleSlots prop1 types1 prop2 types2
      let unifiableAligns := alignments.filter (·.unifiable)

      if unifiableAligns.size > 0 then
        let totalSlots := prop1.varSlots.size + prop2.varSlots.size
        let score := computeScore alignments totalSlots
        candidates := candidates.push {
          theorem1   := thm1.toString
          theorem2   := thm2.toString
          cluster1Id := c1.id
          cluster2Id := c2.id
          shape1     := c1.shapeKey
          shape2     := c2.shapeKey
          alignments := alignments
          score      := score
        }

  return candidates

-- ════════════════════════════════════════════════════════════════════════════
-- § 5  Preview and execute convolution
-- ════════════════════════════════════════════════════════════════════════════

private def combineShapes (s1 s2 : String) : String :=
  s!"({s1}{s2})"

/-- Preview what the composed pattern would look like. -/
def previewConvolution
    (thm1 : Name) (clusterId1 : Nat)
    (thm2 : Name) (clusterId2 : Nat) : MetaM ComposedPattern := do
  let graph1 ← findComonadsCore thm1
  let graph2 ← findComonadsCore thm2

  let cluster1 ← match graph1.clusters.find? (·.id == clusterId1) with
    | some c => pure c
    | none   => throwError s!"cluster {clusterId1} not found in {thm1}"

  let cluster2 ← match graph2.clusters.find? (·.id == clusterId2) with
    | some c => pure c
    | none   => throwError s!"cluster {clusterId2} not found in {thm2}"

  let ci1 ← getConstInfo thm1
  let ci2 ← getConstInfo thm2
  let val1 ← match ci1.value? with | some e => pure e | none => throwError "no value"
  let val2 ← match ci2.value? with | some e => pure e | none => throwError "no value"
  let entries1 ← Mathlib.Explode.explode val1
  let entries2 ← Mathlib.Explode.explode val2
  let richNodes1 ← buildRichIR entries1
  let richNodes2 ← buildRichIR entries2

  let prop1 := analyzeCluster cluster1
  let prop2 := analyzeCluster cluster2
  let types1 := getVarSlotTypes prop1 richNodes1
  let types2 := getVarSlotTypes prop2 richNodes2
  let alignments := findCompatibleSlots prop1 types1 prop2 types2

  let unifiedAligns := alignments.filter (·.unifiable)
  let paramCount := prop1.varSlots.size + prop2.varSlots.size - unifiedAligns.size

  let combinedShape := combineShapes cluster1.shapeKey cluster2.shapeKey

  let paramDecls := (Array.range paramCount).toList.map fun i =>
    s!"(p{i} : _)"
  let paramStr := " ".intercalate paramDecls
  let leanCommand := s!"private def composed_pattern {paramStr} : _ := sorry"

  let newGraphId := hashPPExpr (graph1.graphId ++ graph2.graphId ++ combinedShape)

  return {
    name          := "composed_pattern"
    source1       := (thm1.toString, clusterId1)
    source2       := (thm2.toString, clusterId2)
    alignments    := alignments
    combinedShape := combinedShape
    paramCount    := paramCount
    leanCommand   := leanCommand
    newGraphId    := newGraphId
  }

/-- Execute convolution and generate the composed pattern. -/
def executeConvolution
    (thm1 : Name) (clusterId1 : Nat)
    (thm2 : Name) (clusterId2 : Nat)
    (name : String) : MetaM ComposedPattern := do
  let preview ← previewConvolution thm1 clusterId1 thm2 clusterId2

  let paramDecls := (Array.range preview.paramCount).toList.map fun i =>
    s!"(p{i} : _)"
  let paramStr := " ".intercalate paramDecls
  let leanCommand := s!"private def {name} {paramStr} : _ := sorry"
  let newGraphId := hashPPExpr (preview.newGraphId ++ name)

  return { preview with
    name := name
    leanCommand := leanCommand
    newGraphId := newGraphId
  }

-- ════════════════════════════════════════════════════════════════════════════
-- § 6  JSON serialization
-- ════════════════════════════════════════════════════════════════════════════

private def floatToJson (f : Float) : Lean.Json :=
  if f == 0.0 then .num ⟨0, 0⟩
  else .num { mantissa := Int.ofNat (f * 1000000.0).round.toUInt64.toNat
              exponent := 6 }

private def alignmentToJson (a : SlotAlignment) : Lean.Json :=
  .mkObj [ ("pos1",      .num ⟨Int.ofNat a.pos1, 0⟩)
         , ("pos2",      .num ⟨Int.ofNat a.pos2, 0⟩)
         , ("type1",     .str a.typeStr1)
         , ("type2",     .str a.typeStr2)
         , ("unifiable", .bool a.unifiable) ]

private def candidateToJson (c : ConvolutionCandidate) : Lean.Json :=
  .mkObj [ ("theorem1",   .str c.theorem1)
         , ("theorem2",   .str c.theorem2)
         , ("cluster1",   .num ⟨Int.ofNat c.cluster1Id, 0⟩)
         , ("cluster2",   .num ⟨Int.ofNat c.cluster2Id, 0⟩)
         , ("shape1",     .str c.shape1)
         , ("shape2",     .str c.shape2)
         , ("alignments", .arr (c.alignments.map alignmentToJson))
         , ("score",      floatToJson c.score) ]

def candidatesToJson (candidates : Array ConvolutionCandidate) : Lean.Json :=
  .mkObj [ ("schema", .str "candidates/1")
         , ("pairs",  .arr (candidates.map candidateToJson)) ]

def composedToJson (p : ComposedPattern) : Lean.Json :=
  .mkObj [ ("schema",       .str "preview/1")
         , ("name",         .str p.name)
         , ("source1",      .mkObj [("theorem", .str p.source1.1),
                                    ("clusterId", .num ⟨Int.ofNat p.source1.2, 0⟩)])
         , ("source2",      .mkObj [("theorem", .str p.source2.1),
                                    ("clusterId", .num ⟨Int.ofNat p.source2.2, 0⟩)])
         , ("alignments",   .arr (p.alignments.map alignmentToJson))
         , ("combinedShape",.str p.combinedShape)
         , ("paramCount",   .num ⟨Int.ofNat p.paramCount, 0⟩)
         , ("leanCommand",  .str p.leanCommand)
         , ("newGraphId",   .str p.newGraphId) ]

def composeResultToJson (p : ComposedPattern) : Lean.Json :=
  .mkObj [ ("schema",       .str "compose/1")
         , ("name",         .str p.name)
         , ("source1",      .mkObj [("theorem", .str p.source1.1),
                                    ("clusterId", .num ⟨Int.ofNat p.source1.2, 0⟩)])
         , ("source2",      .mkObj [("theorem", .str p.source2.1),
                                    ("clusterId", .num ⟨Int.ofNat p.source2.2, 0⟩)])
         , ("alignments",   .arr (p.alignments.map alignmentToJson))
         , ("combinedShape",.str p.combinedShape)
         , ("paramCount",   .num ⟨Int.ofNat p.paramCount, 0⟩)
         , ("leanCommand",  .str p.leanCommand)
         , ("newGraphId",   .str p.newGraphId) ]

end Convolution

-- ════════════════════════════════════════════════════════════════════════════
-- § 7  Commands
-- ════════════════════════════════════════════════════════════════════════════

syntax (name := patternComposeCmd) "#patternCompose" ident ident : command

@[command_elab patternComposeCmd]
def elabPatternCompose : CommandElab := fun stx => do
  let thm1 := stx[1].getId
  let thm2 := stx[2].getId
  let candidates ← liftTermElabM (Convolution.findCandidates thm1 thm2)
  if candidates.isEmpty then
    logInfo m!"No convolution candidates found between {thm1} and {thm2}"
  else
    logInfo m!"{(Convolution.candidatesToJson candidates).compress}"

syntax (name := patternPreviewCmd) "#patternPreview" ident num ident num : command

@[command_elab patternPreviewCmd]
def elabPatternPreview : CommandElab := fun stx => do
  let thm1 := stx[1].getId
  let c1 := stx[2].toNat
  let thm2 := stx[3].getId
  let c2 := stx[4].toNat
  let preview ← liftTermElabM (Convolution.previewConvolution thm1 c1 thm2 c2)
  logInfo m!"{(Convolution.composedToJson preview).compress}"

syntax (name := patternExecuteCmd) "#patternExecute" ident num ident num "as" str : command

@[command_elab patternExecuteCmd]
def elabPatternExecute : CommandElab := fun stx => do
  let thm1 := stx[1].getId
  let c1 := stx[2].toNat
  let thm2 := stx[3].getId
  let c2 := stx[4].toNat
  let name := stx[6].isStrLit?.getD "composed"
  let result ← liftTermElabM (Convolution.executeConvolution thm1 c1 thm2 c2 name)
  logInfo m!"{(Convolution.composeResultToJson result).compress}"

-- (Test section dropped on port — see ComonadFinder § 13 note;
-- the `#patternCompose` elaborator may also need updating for
-- Lean 4 v4.30, which is part of the deferred test-restoration.)
