import std/unittest
import std/math
import std/sequtils
import std/sugar
import std/random

import constantine/math/arithmetic
import constantine/math/io/io_fields
import constantine/math/io/io_bigints
import constantine/serialization/codecs

import poseidon2/types
import poseidon2/io
import poseidon2/compress
import poseidon2/merkle

import ./fuzzing
import ./reference

suite "merkle root":

  const isBottomLayer = 1
  const isOddNode = 2

  test "merkle root of field elements":
    let m = 17
    let n = 2^m
    var xs: seq[F]
    for i in 1..n:
      xs.add( toF(i) )

    let root = Merkle.digest(xs)
    check root.toHex(littleEndian) == "0x593e01f200cb1aee4e75fe2a9206abc3abd2a1216ab75f1061965e97371e8623"

  test "merkle root of even elements":
    let elements = toSeq(1..4).mapIt(toF(it))
    let expected = compress(
      compress(1.toF, 2.toF, key = isBottomLayer.toF),
      compress(3.toF, 4.toF, key = isBottomLayer.toF),
    )
    check bool(Merkle.digest(elements) == expected)

  test "merkle root of odd elements":
    let elements = toSeq(1..3).mapIt(toF(it))
    let expected = compress(
      compress(1.toF, 2.toF, key = isBottomLayer.toF),
      compress(3.toF, 0.toF, key = (isBottomLayer + isOddNode).toF)
    )
    check bool(Merkle.digest(elements) == expected)

  test "data ending with 0 differs from padded data":
    let a = toSeq(1..3).mapIt(it.toF)
    let b = a & @[0.toF]
    check not bool(Merkle.digest(a) == Merkle.digest(b))

  test "merkle root of single element does not equal the element":
    check not bool(Merkle.digest([1.toF]) == 1.toF)

  test "merkle root differs from merkle root of merkle root":
    let a = 1.toF
    let b = 2.toF
    check not bool(Merkle.digest([a, b]) == Merkle.digest([Merkle.digest([a, b])]))

  test "merkle root of bytes":
    let bytes = toSeq 1'u8..80'u8
    let root = Merkle.digest(bytes)
    check root.toHex(littleEndian) == "0x40989b63104f39e3331767883381085bcfc46e2202679123371f1ffe53521b16"

  test "merkle root of bytes converted to bytes":
    let bytes = toSeq 1'u8..80'u8
    let rootAsBytes = Merkle.digest(bytes).toBytes()
    check rootAsBytes.toHex == "0x40989b63104f39e3331767883381085bcfc46e2202679123371f1ffe53521b16"

  test "merkle root of empty sequence of elements":
    let empty = seq[F].default
    expect Exception:
      discard Merkle.digest(empty)

  test "merkle root of empty sequency of bytes":
    # merkle root of empty sequence of bytes is uniquely defined through padding
    let empty = seq[byte].default
    check Merkle.digest(empty).toBytes.toHex == "0xcc8da1d157900e611b89e258d95450e707f4f9eec169422d7c26aba54f803c08"

suite "merkle root test vectors":

  test "field elements":

    const expected = [
      "03725399183367945352080398854175773551921581713520486387171444673504688049612",
      "01200363431219114414119550523646199479423259809629365937886754089111624051137",
      "03290849705974295885356475812949977947719075082723205888372484144436587857608",
      "13320207757774496338093403190247235704739125936593833502280725662388374071598",
      "08797512419619623354301868676697660408674060215007182352266699867257089555918",
      "18775477084402365457164097678179278781831144424754858446675853467823831685773",
      "05618925189910878733331114718351946258062126524237918283575709000172677619746",
      "02468800965850777178862816556314777665879714166580718063606541428124645523179",
      "21298151378974529563336714932691053445808366094493516245922754590652455011873",
      "06347201706526769486134739281293106492815906958149583545680894660915362714427",
      "16618343244443304744112561205464374214211409693479461100629842689638295628832",
      "01709215889546441870644088689718294006086901651959914593564225542798292183537",
      "14773444040872759628462450230133287467645269839387662282102026343061616552943",
      "06108821195988570207520576300404160824923407981901092560371182595716182192554",
      "05022687626412485912451894176261697772833573996468387114965205048748116063384",
      "16980404023653621595001952565345111473785834266619087676781785983708180485792",
      "15853530753614161257415924211864771913403724603890733979349732983529985595273",
      "15031913754244602537901192245353232049056918360300952523716286915084019590221",
      "16862094283091684687208785360745140350115474186985529726891156327026129253080",
      "16492021048693292688486386321061780771759395012813093173702769376227716192997",
      "20278966451219169781380490490108162212827856988184513236302841504016649059914",
      "07432269657144814551290369489541295309360618460337125447714892164714595690148",
      "10983974081267177960524156069527964183879935469683463758001589348825879472639",
      "07651583156777324186897885604441033008761341166782113340300149936063411446625",
      "07590084480733089221411733320811992183469656506657093732150221762955528201604",
      "00590603699839296952039093394825098803770404112267844315790208005762987394025",
      "06359371005747954026039452744396819852623439843092320469786606103532367578946",
      "13746527612045411798909927669111702188976071788637707551527016974970749581239",
      "15119908003571853695118996260713137032585853691267164461203689670845264410241",
      "17682380662262842717089506798091895756812609600205063094408086964114510882038",
      "21269428751047382014946948997891826023229683848096455558245783700971446452460",
      "18760257053115407642276950856798213638198967648053072974895207590929370321671",
      "13634513843727898927772335038794285458655559239614377841772627165816947638955",
      "10006312780138777086841229248944999114254866552257803132887344279098482458747",
      "01843241079155742770064553335351435848667557880058014673413647601991317785477",
      "11272460975973875704457193928858301545715516859429881611780245005127850725636",
      "04582167569197996558082729354402407015968305791271215294732707709246372235973",
      "21370100954705769476406168562536690546315894860847236971113013796274873562782",
      "11529800546649224443767229283005832742131266548011774867883901788819175729967",
      "01444081399852704913168826065775287766195189673591523318210191867140358619452",
    ]

    for n in 1..40:
      let input = collect(newSeq, (for i in 1..n: i.toF))
      let root = Merkle.digest(input)
      check root.toDecimal == expected[n-1]

  test "byte sequences":

    const expected = [
      "03725399183367945352080398854175773551921581713520486387171444673504688049612",
      "21784172376012850120488664011679965109942973874126334661725579904088599292772",
      "07152079943212903152545420874069167733062870609755361597526511154461589187355",
      "08141473468355887323169573159427773162133489060487824719557688349428817589204",
      "18849072884479535235752900319646951743421793496758367608904079738623609476991",
      "01035096315938703794310832272083769709349075272522073827895849963478820268491",
      "13445844440885991019137991459933024269053128830981235694187117956602081299688",
      "01590178538056532973070549921963393506524495989294396711240545074117105763737",
      "20309296464966529825402717045955756083491741740781344827838165328540049733679",
      "00126874371085135214945275787718436302709445321127917734525681740069234612287",
      "07786225353194809237480895721268481020740747794118503424204769381994616595569",
      "10610825450875569929871347555810699348068265672369268032037661130733542193887",
      "00071049007785349661407127486289880418543135913647515254842487059628387717609",
      "09277045393414933079438810285584246716893106513327068553120404169267467939019",
      "07982633314069083175183001875148208854326096823785299812521826177919623165496",
      "04939112189475757820362273909178316495851390181049394853623408416533024237576",
      "11006371764505320440420887144360891084304582358134826465620068888693583722497",
      "10881763829906719104039450195610605062046978887183595820306836165133993662959",
      "08802003101472014948343991097326088351898737134336033964330381996709551314119",
      "15636435222554850008213199017027690196983900581613505014605413102138485986245",
      "07335657549016427934228374286432235373964786888487118874028129821852808882463",
      "03195093225419946921825396953874176599941014011418264404491245512267011194152",
      "21560635004945011828405021563208481492862884537481032540645770248759652670216",
      "12641524701910776242674631919516987960769706525915138736344357088285908182252",
      "08675566246253481781842086337026282437389565060406594089426619279151156640731",
      "05676872618423617878730087437442286697922050696270078476033728075327041035735",
      "16680961928897795533088492217095293519890407760847734324403615466542137196978",
      "09049220677560826807459659002105278977433365116260239159032160443735579451094",
      "21608318712162283302901938465434609794460953178533966662576890871495390853513",
      "15050697512091214875071760411984429947277434512680774490434388628001602806946",
      "04332025221097145854023632834862365808876796053617527838123176962567507647936",
      "12299074978003881118784365283411816695529572541077146649102337432570242500854",
      "10783003198936970298455299876751834916729312690493205407592190957968056038658",
      "02603124601761869438414850819289748229671771473053126931734122942024553197832",
      "07110864127516205437690550396435710305236648267482065174460237774365233651950",
      "08775867381899869809760736769992319818197850252740245847574699244299061615645",
      "00850175513812112390640199871976516255016499379401357704447229391967903811913",
      "18781216421511020669819637004471959594470202972018133188711575252428764314600",
      "06368706098954594854304368410137650649779615071696538040253930550477917178402",
      "09019570936956533560062914106955526624646648228064181779449758922204503303704",
      "05440887168154527956624914122727365035123822752834072410988856388113119571470",
      "20550068893651968253698431875989392535090061052279579955555002682668390985213",
      "15737389942660057327683444756111102130453778067838029605808425612712027337551",
      "18815402418895765221293726188836046757016812774570433066993070555279333481533",
      "11682323118384360956769167276674716180735339639693415278211786113335964794851",
      "06935660702899622014808820492490754833483949376170647030039377716372799733842",
      "21810909209286088906970343552757991729652062252952401401820661936760194844461",
      "04342808204475053897845478423167674781330792605324084630256784374136346846489",
      "15374402795399778784660786340596814682805690545616806818642115141863066723874",
      "13447787026254285532225826718566842976436581885719486841551519403729861198889",
      "12936970407640825952749594792510039878330172708509401147261022071094145009815",
      "09164059512510832875663760138394876191213584366639343544153038610001712299562",
      "01413614220784930065401416340789241055974640105288991358124042204900619845043",
      "04536794469882872029454805585069444556045288475866146536311433334366441590422",
      "04065111163193508717294204957954496537740683600048901299319693465444044780399",
      "07574840121491782797776968017659317905710680028291006386390853521815661119797",
      "06184384950445868673596682556762687453458927278845452900030573959958522289156",
      "15839307597459993108945519517841264511212618857558276342907222864813120675506",
      "11207712115697791710896047063384302468090264935915626542301613733795025526404",
      "12232986987162313228928134416576881145873864942675997264125317200642330847158",
      "10799160640818664650911188873129536812944809901077658227097120255766699516430",
      "00554553831389315183486861104334894017052073737212863865655082667512560253517",
      "13772907771367224435776564966368426336775487238426796618778567296717838253898",
      "08850067274802611078644005969765667843461068373958906729693103015368591725949",
      "20058234434689951307215284422042343876303897330000168535619938965778013422946",
      "04630140766537898447420609228017596645008053776260989301156385378345683670482",
      "17736653050758446236565266013480353492407731491080899492469529050713235546697",
      "05378965009312241887537092288550907102922450643751682252500303378152366939458",
      "01295499824483568650260107404264570051256567099921276873258070841350453934067",
      "07583535500986628910603010605606141497422970011939147970590627698261229620371",
      "11966184222148654156965514341908034850722389372357084477913489948173230763118",
      "09248921860745338533460221363471205215550148340610364237016637502708485523921",
      "08523227948351872181200239631998481886943486401829210258780254835047235663156",
      "04899178581629450481188598886000169627699337799896231567527097023198715671176",
      "10177511949352582811609663265581377778190428048679249238548931306303103613925",
      "08639576565773586801169549841316728276022076647431203118061123177280288490728",
      "07258854121558679587397850669889059677710821136114798324609764439794142871434",
      "13256806563302067064811128237965864768596645306266796341089300636189766949669",
      "07809487468511092855854823943165286878075560110185627869354817357553972592025",
      "11200627159277717358031252616982173140143711715362282993454547889930173794425",
      "09999155747219012026170409690829745881511413685936981652054886005773929519168",
    ]

    for n in 0..80:
      let input = collect(newSeq, (for i in 1..n: byte(i)))
      let root = Merkle.digest(input)
      check root.toDecimal == expected[n]

suite "merkle root fuzzing (seed: " & $fuzzing.seed & ")":

  test "merkle root algorithm matches reference implementation":

    proc checkInput(input: seq[byte]) =
      let expected = reference.merkleRoot(input)
      check bool(Merkle.digest(input) == expected)

    # a couple of tests with small input
    for _ in 0..<1000:
      let len = rand(1024)
      let input = newSeqwith(len, byte.rand())
      checkInput(input)

    # one test with larger input
    let len = rand(1024 * 1024)
    let input = newSeqwith(len, byte.rand())
    checkInput(input)