"""
- Authors: Nissi, Jada, Rachel, Kaitlyn, Kitty, Peter Mawhorter
- Consulted: Peter Mawhorter
- Date: 2022-12-5
- Purpose: Tests for exploration analysis.
"""

from .. import analysis, journal, core, base

#Part of Rachel's journal:
JOURNAL = """
S Start_room::Start
  zz Starting_region
  A gain attack
  o right
  o left
xt left coin_room sr
  o up above_coin_room fall
    q _wall_kick_jumps  # TODO: '?' syntax
t sr
xt right platforms rp
  o up
  o right
xt up dangerplat ud
"""

BABY_JOURNAL = """
S Start
  A gain jump
  A gain attack
  n button check
  zz Wilds
  o up
    q _flight
  o left
xt left left_nook right
  a geo_rock
    At gain geo*15
    At deactivate
  o up
    q _tall_narrow
t right
  o right
    q attack
"""

BABY_JOURNAL_2 = """
S Start
  A gain jump
  A gain attack
  n button check
  zz Wilds
  o up
    q _flight
  o left
xt left left_nook right
"""

REVISITS_JOURNAL = """
S A
x right B left
x down C up
x left D right           # 4 decisions ABCD
r up A down              # revisit #1 for A
t right                  # revisit #1 for B
r down_left D up_right   # revisit #1 for D
t right                  # revisit #1 for C
t up                     # revisit #2 for A
t left                   # revisit #2 for B
t down                   # revisit #2 for C

# now we have: 2 revisits each for A, B, and C, and 1 revisit for D
"""

ACTIONS_JOURNAL = """
S A
  oa alpha
x right B left
  oa beta
t left
  ta alpha
  oa gamma
x down C up
  a delta
  a epsilon
x right D left
w
"""


def test_CreateExploration() -> None:
    """
    Simple test to make sure we can create an exploration object for
    other tests in this file.
    """
    ex = journal.convertJournal(JOURNAL)
    assert len(ex) == 6
    assert ex.getActiveDecisions(0) == set()
    assert ex.getActiveDecisions(1) == {0}

    assert bool(ex.getSituation(1).graph) is True


def test_countActionsAtDecision() -> None:
    """
    Test to make sure the number of actions in the whole graph
    at each step is accurate.
    """
    ex = journal.convertJournal(BABY_JOURNAL)
    assert [
        analysis.totalActions(ex, i)
        for i in range(len(ex))
    ] == [0, 0, 1, 1, 1]

    now = ex.getSituation()
    graph = now.graph
    startID = graph.resolveDecision("Start")
    nookID = graph.resolveDecision("left_nook")
    assert analysis.actionCount(ex, len(ex) - 1, startID) == 0
    assert analysis.actionCount(ex, len(ex) - 1, nookID) == 1

    ex2 = journal.convertJournal(ACTIONS_JOURNAL)
    assert [
        analysis.totalActions(ex2, i)
        for i in range(len(ex2))
    ] == [0, 1, 2, 2, 3, 4, 5, 5, 5, 5]
    assert [
        analysis.meanActions(ex2, i)
        for i in range(len(ex2))
    ] == [0.0, 1/2, 2/2, 2/2, 3/3, 4/3, 5/3, 5/4, 5/4, 5/4]

    for step in range(len(ex2)):
        assert (
            analysis.totalActions(ex2, step) /
            analysis.totalDecisionsSoFar(ex2, step)
         == analysis.meanActions(ex2, step)
        )


def test_count_actions() -> None:
    """
    Tests the 'actionCount' metric and its total/mean combiner values.
    """
    ex = journal.convertJournal(BABY_JOURNAL)
    ex2 = journal.convertJournal(ACTIONS_JOURNAL)

    assert analysis.actionCount(ex, 1, 0) == 0
    assert analysis.actionCount(ex, 2, 1) == 0


def test_describeProgress() -> None:
    """
    Tests the `describeProgress` function.
    """
    e1 = journal.convertJournal(BABY_JOURNAL)
    e2 = journal.convertJournal(BABY_JOURNAL_2)
    e3 = journal.convertJournal(REVISITS_JOURNAL)

    description = analysis.describeProgress(e1)
    assert description == """\
Start of the exploration
Start exploring domain main at 0 (Start)
  Gained capability 'attack'
  Gained capability 'jump'
At decision 0 (Start)
  In region Wilds
  There are transitions:
    left to unconfirmed
    up to unconfirmed; requires _flight
  1 note(s) at this step
Explore left from decision 0 (Start) to 2 (now Wilds::left_nook)
At decision 2 (left_nook)
  There are transitions:
    right to 0 (Start)
  There are actions:
    geo_rock
Do action geo_rock
  Gained 15 geo(s)
Take right from decision 2 (left_nook) to 0 (Start)
At decision 0 (Start)
  There are transitions:
    left to 2 (left_nook)
    right to unconfirmed; requires attack
    up to unconfirmed; requires _flight
Waiting for another action...
End of the exploration.
"""

    description2 = analysis.describeProgress(e2)
    assert description2 == """\
Start of the exploration
Start exploring domain main at 0 (Start)
  Gained capability 'attack'
  Gained capability 'jump'
At decision 0 (Start)
  In region Wilds
  There are transitions:
    left to unconfirmed
    up to unconfirmed; requires _flight
  1 note(s) at this step
Explore left from decision 0 (Start) to 2 (now Wilds::left_nook)
At decision 2 (left_nook)
  There are transitions:
    right to 0 (Start)
Waiting for another action...
End of the exploration.
"""

    description3 = analysis.describeProgress(e3)
    assert description3 == """\
Start of the exploration
Start exploring domain main at 0 (A)
At decision 0 (A)
  There are transitions:
    right to unconfirmed
Explore right from decision 0 (A) to 1 (now B)
At decision 1 (B)
  There are transitions:
    down to unconfirmed
    left to 0 (A)
Explore down from decision 1 (B) to 2 (now C)
At decision 2 (C)
  There are transitions:
    left to unconfirmed
    up to 1 (B)
Explore left from decision 2 (C) to 3 (now D)
At decision 3 (D)
  There are transitions:
    right to 2 (C)
    up to 0 (A)
Take up from decision 3 (D) to 0 (A)
At decision 0 (A)
  There are transitions:
    down to 3 (D)
    right to 1 (B)
Take right from decision 0 (A) to 1 (B)
At decision 1 (B)
  There are transitions:
    down to 2 (C)
    down_left to 3 (D)
    left to 0 (A)
Take down_left from decision 1 (B) to 3 (D)
At decision 3 (D)
  There are transitions:
    right to 2 (C)
    up to 0 (A)
    up_right to 1 (B)
Take right from decision 3 (D) to 2 (C)
At decision 2 (C)
  There are transitions:
    left to 3 (D)
    up to 1 (B)
Take up from decision 2 (C) to 1 (B)
At decision 1 (B)
  There are transitions:
    down to 2 (C)
    down_left to 3 (D)
    left to 0 (A)
Take left from decision 1 (B) to 0 (A)
At decision 0 (A)
  There are transitions:
    down to 3 (D)
    right to 1 (B)
Take down from decision 0 (A) to 3 (D)
At decision 3 (D)
  There are transitions:
    right to 2 (C)
    up to 0 (A)
    up_right to 1 (B)
Waiting for another action...
End of the exploration.
"""


def test_unexploredBranches() -> None:
    """
    Tests the `unexploredBranches` and related count functions.
    """
    ex = journal.convertJournal(JOURNAL)
    assert analysis.unexploredBranches(ex.getSituation(0).graph) == []
    g1 = ex.getSituation(1).graph
    assert g1.destinationsFrom(0) == {
        'right': 1,
        'left': 2
    }
    assert g1.nameFor(1) == '_u.0'
    assert g1.nameFor(2) == '_u.1'
    assert analysis.unexploredBranches(ex.getSituation(1).graph) == [
        (0, 'right'),
        (0, 'left'),
    ]
    assert analysis.unexploredBranches(ex.getSituation(2).graph) == [
        (0, 'right'),
        (2, 'up'),
    ]
    assert analysis.unexploredBranches(ex.getSituation(3).graph) == [
        (0, 'right'),
        (2, 'up'),
    ]
    g4 = ex.getSituation(4).graph
    assert g4.namesListing(set(g4.nodes)) == """\
  0 (Start_room::Start)
  1 (Start_room::platforms)
  2 (Start_room::coin_room)
  3 (Start_room::above_coin_room)
  4 (_u.3)
  5 (_u.4)
"""
    assert analysis.unexploredBranches(ex.getSituation(4).graph) == [
        (1, 'up'),
        (1, 'right'),
        (2, 'up'),
    ]
    assert analysis.unexploredBranches(ex.getSituation(5).graph) == [
        (1, 'right'),
        (2, 'up'),
    ]
    allPerStep = [
        analysis.unexploredBranchCount(ex, step)
        for step in range(len(ex))
    ]
    traversablePerStep = [
        analysis.traversableUnexploredCount(ex, step)
        for step in range(len(ex))
    ]
    assert allPerStep == [0, 2, 2, 2, 3, 2]
    assert traversablePerStep == [0, 2, 1, 1, 2, 1]
    # TODO


def test_countBranches() -> None:
    """
    Tests the `countBranches` function.
    """
    ex = journal.convertJournal(BABY_JOURNAL)
    ex2 = journal.convertJournal(BABY_JOURNAL_2)

    # Note: as of v0.6, we can index an exploration to get a Situation,
    # and most analysis functions want Situations as input

    first = ex[0]
    second = ex[1]
    third = ex[2]
    fourth = ex[3]
    fifth = ex[4]

    ex2first = ex2[0]
    ex2second = ex2[1]
    ex2third = ex2[2]

    assert first.graph.namesListing(set(first.graph.nodes)) == """\
  0 (Start)
"""
    assert analysis.meanBranches(ex, 0) == 0
    assert analysis.meanBranches(ex, 1) == 2
    assert analysis.meanBranches(ex, 2) == 1.5
    assert analysis.meanBranches(ex, 3) == 2
    assert analysis.meanBranches(ex, 4) == 2.5

    startID = fifth.graph.resolveDecision("Start")
    nookID = fifth.graph.resolveDecision("left_nook")
    assert analysis.branches(ex, 1, startID) == 2
    assert analysis.branches(ex, 2, startID) == 2
    assert analysis.branches(ex, 2, nookID) == 1
    assert analysis.branches(ex, 4, startID) == 3
    assert analysis.branches(ex, 4, nookID) == 2

    assert analysis.meanBranches(ex2, 0) == 0
    assert analysis.meanBranches(ex2, 1) == 2
    assert analysis.meanBranches(ex2, 2) == 1.5


def test_revisits() -> None:
    """
    Tests the `arrivals`, `revisits`, and related total/mean/median
    functions.
    """

    babyExp = journal.convertJournal(BABY_JOURNAL)
    graph = babyExp.getSituation().graph
    startID = graph.resolveDecision("Start")
    nookID = graph.resolveDecision("left_nook")
    assert analysis.stepsVisited(babyExp) == {
        startID: [1, 4],
        nookID: [2, 3]
    }
    assert analysis.arrivals(babyExp, startID) == 2
    assert analysis.revisits(babyExp, startID) == 1
    assert analysis.arrivals(babyExp, nookID) == 1
    assert analysis.revisits(babyExp, nookID) == 0
    assert analysis.arrivals(babyExp, 28309823) == 0
    assert analysis.revisits(babyExp, 28309823) == 0

    fullExp = journal.convertJournal(JOURNAL)
    graph = fullExp.getSituation().graph
    startID = graph.resolveDecision("Start")
    coinRoomID = graph.resolveDecision("coin_room")
    platformsID = graph.resolveDecision("platforms")
    dangerplatID = graph.resolveDecision("dangerplat")
    assert analysis.revisits(fullExp, startID) == 1
    assert analysis.revisits(fullExp, coinRoomID) == 0
    assert analysis.revisits(fullExp, platformsID) == 0
    assert analysis.revisits(fullExp, dangerplatID) == 0

    revExp = journal.convertJournal(REVISITS_JOURNAL)
    graph = revExp.getSituation().graph
    aID = graph.resolveDecision('A')
    bID = graph.resolveDecision('B')
    cID = graph.resolveDecision('C')
    dID = graph.resolveDecision('D')
    assert analysis.revisits(revExp, aID) == 2
    assert analysis.revisits(revExp, bID) == 2
    assert analysis.revisits(revExp, cID) == 1
    assert analysis.revisits(revExp, dID) == 2

    assert analysis.totalRevisits(babyExp) == 1
    assert analysis.meanRevisits(babyExp) == 1/5
    assert analysis.medianRevisits(babyExp) == 0
    assert analysis.totalRevisits(fullExp) == 1
    assert analysis.meanRevisits(fullExp) == 1/6
    assert analysis.medianRevisits(fullExp) == 0
    assert analysis.totalRevisits(revExp) == 7
    assert analysis.meanRevisits(revExp) == 7/4
    assert analysis.medianRevisits(revExp) == 2.0

    blankExp = core.DiscreteExploration()
    assert analysis.totalRevisits(blankExp) == 0
    assert analysis.meanRevisits(blankExp) is None
    assert analysis.medianRevisits(blankExp) is None

    soloExp = core.DiscreteExploration()
    soloExp.start("start")
    assert analysis.totalRevisits(soloExp) == 0
    assert analysis.meanRevisits(soloExp) == 0
    assert analysis.medianRevisits(soloExp) == 0

    doubleExp = core.DiscreteExploration()
    doubleExp.start("start")
    doubleExp.observe("start", "up")
    doubleExp.explore("up", "above", "down")
    assert analysis.totalRevisits(doubleExp) == 0
    assert analysis.meanRevisits(doubleExp) == 0
    assert analysis.medianRevisits(doubleExp) == 0


def test_makeFractionCombiner() -> None:
    """
    Tests `makeFractionCombiner`.
    """
    c = analysis.makeFractionCombiner(lambda k, x: x > 2, lambda k, x: x < 0)
    d = {
        'a': 1,
        'b': -2,
        'c': 5,
        'd': 3,
        'e': 2.1
    }
    assert c(d) == 0.75  # type:ignore
    c = analysis.makeFractionCombiner(lambda k, x: x < 0, lambda k, x: x > 0)
    assert c(d) == 1.0  # type:ignore
    c = analysis.makeFractionCombiner(lambda k, x: x < 0)
    assert c(d) == 0.2  # type:ignore
    c = analysis.makeFractionCombiner(lambda k, x: x > 2)
    assert c(d) == 0.6  # type:ignore
    c = analysis.makeFractionCombiner(lambda k, x: x % 3 == 7)
    assert c(d) == 0.0  # type:ignore

def test_cycleBasis() -> None:
    """
    Tests `cycleBasis`.
    """
    # Test using a custom several-cycle decision graph:
    g = core.DecisionGraph()
    g.addDecision('A')  # 0
    g.addDecision('B')  # 1
    g.addDecision('C')  # 2
    g.addDecision('D')  # 3
    g.addDecision('E')  # 4
    g.addDecision('F')  # 5
    g.addDecision('G')  # 6
    g.addDecision('H')  # 7
    g.addDecision('I')  # 8
    assert analysis.cycleBasis(g) == []
    g.addTransition('A', 'left', 'B', 'right')
    g.addTransition('B', 'down', 'C', 'upRight')
    assert analysis.cycleBasis(g) == []
    g.addTransition('A', 'down', 'C', 'upLeft')  # triangle
    assert analysis.cycleBasis(g) == [[0, 1, 2]]
    g.addTransition('C', 'down', 'D', 'up')
    g.addTransition('D', 'down', 'E', 'up')  # D is in-between
    assert analysis.cycleBasis(g) == [[0, 1, 2]]
    g.addTransition('E', 'downLeft', 'F', 'up')
    g.addTransition('F', 'right', 'G', 'left')
    g.addTransition('G', 'right', 'H', 'left')
    assert analysis.cycleBasis(g) == [[0, 1, 2]]
    g.addTransition('H', 'up', 'E', 'downRight')  # 4-node cycle
    assert {frozenset(c) for c in analysis.cycleBasis(g)} == {
        frozenset({0, 1, 2}),
        frozenset({4, 5, 6, 7})
    }
    g.addTransition('G', 'down', 'I', 'up')  # branches off
    assert {frozenset(c) for c in analysis.cycleBasis(g)} == {
        frozenset({0, 1, 2}),
        frozenset({4, 5, 6, 7})
    }
    g.setTransitionRequirement('C', 'down', base.ReqCapability('fly'))
    g.setTransitionRequirement('E', 'up', base.ReqCapability('fly'))
    assert {frozenset(c) for c in analysis.cycleBasis(g)} == {
        frozenset({0, 1, 2}),
        frozenset({4, 5, 6, 7})
    }

    # Test using snippet from problematic journal
    cotmStart = """\
# Setup
P softerChecks on

= enemy which common level rare rareLevel [
  gd enemy:{which}
  oa farm{which}{common}
    e { <{level}>luck{ gain {common}*1 }{ } }
  oa farm{which}{rare}
    e { <{rareLevel}>luck{ gain {rare}*1 }{ } }
]

= farm from what [
  ta farm{from}{what}%s
]

= cardEnemy which card level otherDrop otherLevel [
  gd enemy:{which}
  oa farmCard{which}{card}
    e { <{level}>luck{ gain {card}; deactivate }{ } }
  oa farm{which}{otherDrop}
    e { <{otherLevel}>luck{ gain {otherDrop}*1 }{ } }
]

= gotCard from which [
  ta farmCard{from}{which}%s
]

= subweapon which [
  oa take{which}
    e { lose Dagger; lose HolyWater; lose Boomerang; lose Axe; gain {which} }
]

= getItem what [
  a get{what}
    At { gain {what}*1; deactivate }
]

# Start
S ConfrontationSite::main
  zz Catacombs
  o left
  o down
  m door:sealed
  o right
    qb door:unsealed
  @
    xt down RubbleTowerShaft::shaft up
      gd suspended
  @ @
  a meetDracula
    gt triggered
    At follow down

n at: RubbleTowerShaft::shaft
  o fall bottom flyUp
    qr ?flight
  !a fallDown
    q !?flight
    gt triggered
    At follow fall
  n at: RubbleTowerShaft::bottom

x down RubbleTowerBase::top up
  > cardEnemy SkeletonBomber Salamander 1 Potion 6
    n [ Real chances are 20% / 1%; this gives 25% / 0.78% ]
  > gotCard SkeletonBomber Salamander
    n [ Literally first enemy I touched; super lucky! ]
x down mid up
  > cardEnemy SkeletonBomber Salamander 1 Potion 6
  o leftCrack
    q ?airSlide
x down bottom up
  > cardEnemy SkeletonBomber Salamander 1 Potion 6

x right CatacombsEntrance::top left
  > cardEnemy SkeletonBomber Salamander 1 Potion 6
  o upLeft
    q ?higherJump
  o upRight
    q ?higherJump
  > enemy Bat WristBand 4 Heart 7
    n [ Real 2.5% / 0.5%; these are 3.125% / 0.39% ]
    """
    exp = journal.convertJournal(cotmStart)
    assert analysis.cycleBasis(exp.getSituation().graph) == []