import std/[strutils, unicode]

## This notational model is based on the 12-tone modal, harmonic system
## standard in Western music. A key concept from this is the key-agnostic
## naming system common to Jazz and Gospel music, an expansion of the
## "Nashville Number" system that refers to pitches by their scale degree and
## "alteration" (for lack of a better word). For example, "flat III,"
## "natural/perfect V."
##
## In this model scale degree names are always counted from the tonic,
## regardless of mode. Alterations are based on how the note differs from the
## note at the same scale degree in the Ionian scale that shares the same
## tonic. We also make a distiction between the key-agnostic scale degree and
## alteration, and the key-specific spelling of the same. For example, consider
## the following scales with named notes and key-agnostic scale degrees:
##
## Key of C major
##
##     C  D  E  F  G  A  B
##     1  2  3  4  5  6  7
##
## Key of C minor (Aeolian)
##
##     C  D  E♭ F  G  A♭ B♭
##     1  2  ♭3 4  5  ♭6 ♭7
##
## Key of A minor
##
##     A  B  C  D  E  F  G
##     1  2  ♭3 4  5  ♭6 ♭7
##
## Key of A major
##
##     A  B  C# D  E  F# G#
##     1  2  3  4  5  6  7
##
## Key of G♭ Locrian
##
##     G♭ A𝄫 B𝄫 C♭ D𝄫 E𝄫 F♭
##     1  ♭2 ♭3 4  ♭5 ♭6 ♭7
##
## Key of G Ionian (major)
##
##     G  A  B  C  D  E  F#
##     1  2  3  4  5  6  7
##
## In these examples, C major and A minor are enharomonic (sharing the same
## pitches), spell all of their pitches the same, but have different flavors of
## scale degrees. G Ionian and Gb Locrian are also enharmonic, but have very
## different spellings of the pitches and flavors of scale degrees. Most people
## would use F# Locrian rather than Gb Locrian, because F# Locrian has the same
## pitch spellings as the familiari relative major of G Ionian. The key point
## for understanding this data model is that it allows the author to use either
## and will correctly transpose, using the correctly named pitches based on the
## mode and spelling choice for the key.
##
## This preservation of choices extends to non-diatonic notes as well. For
## example, in the key of C the notes Ab and G# are enharmonically equivalent,
## but functionally different. Ab is the flat VI, whereas G# is the
## augmented/sharp V. There are situations where an author may prefer either of
## these. For example, as a submediant in a walk-up progression, ## bVI-bVII-bI,
## writing it as Cb, the bVI emphasises the whole-tone pattern of the walk up,
## the relationship to VII. In a different progression, a IV-vdim-vi walk up,
## the same pitch might be thought of as a sharp V approaching the VI,
## emphasizing its role as a continuation of the V tension before the
## resolution to the v. Again, this data model allows us to preserve the
## notation of both Cbmaj7-Db6-Eb2 and Bb-Bdim-Cm7 in the key of Eb, notating
## them in the key of D as Bbmaj7-C6-D2 and A-A#dim-Bbm7 respectively.


type
  Note* {.pure.} = enum A, B, C, D, E, F, G
    ## Notes capture the principle in western harmony of the common scales
    ## having seven distinct diatonic "notes" and allows us to do arithmetic
    ## with them (go up three scale degrees).

  Pitch* = enum Af, A, Bf, B, C, Df, D, Ef, E, F, Gf, G
    ## The 12 chromatic pitches, regardless key. Pitch allows us to assign
    ## unique names and ordinal values to each note of the chromatic scale,
    ## allowing us to do arithmetic with them.

  NoteAlteration* = enum
    naDoubleFlat = -2, naFlat, naNatural, naSharp, naDoubleSharp
    ## The supported alterations to a Note or scale number.

  ScaleDegree* = object
    number: int
    alteration: NoteAlteration
    ## A key-agnostic representation of a pitch, capturing the scale degree
    ## relative to the tonic, and the alteration of that. So, for example:
    ## b3 -> b (alteration) 3 (number)

  SpelledPitch* = object
    note: Note
    alteration: NoteAlteration
    ## A unique spelling of one of the chromatic pitches allowing the stylistic
    ## choice to use different *Notes* to describe a single *Pitch*. For
    ## example, Cb, A𝄪, and B♮ are three distinct ways to spell Pitch.B

  Mode* = enum
    Ionian = 0,
    Dorian,
    Phrygian,
    Lydian,
    Mixolydian,
    Aeolian,
    Locrian

  Key* = object
    tonic*: SpelledPitch
    mode*: Mode

const MajorIntervals = [0, 2, 4, 5, 7, 9, 11]


func `+`*[T: Pitch or Note](val: T, steps: int): T =
  ## Move up or down the diatonic scale or the chromatic scale by the given
  ## number of steps.
  cast[T]((ord(val) + steps) mod (ord(high(T)) + 1))


func `-`*[T: Pitch or Note](val: T, steps: int): T =
  ## Move down the diatonic or chromatic scale by the given number of steps.
  var newOrd = (ord(val) - steps) mod (ord(high(T)) + 1)
  if newOrd < 0: newOrd += 12
  return cast[T](newOrd)


func `-`*(a, b: Note): int =
  ## Find the distance between two notes of the diatonic scale. This always
  ## returns a positive distance, as if `a` was higher in pitch than `b`. For
  ## example, C - D returns +6 (the distance from D3 to C4) rather than -1
  ## (the distance from D4 to C4).
  result = ord(a) - ord(b)
  if result < 0: result += 7


func `-`*(a, b: Pitch): int =
  ## Find the distance between two notes of the chromatic scale. This always
  ## returns a positive distance, as if `a` was higher in pitch than `b`. For
  ## example, C - D returns +10 (the distance from D3 to C4) rather than -2
  ## (the distance from D4 to C4).
  result = ord(a) - ord(b)
  if result < 0: result += 12


func `$`*(pitch: Pitch): string =
  ["A♭", "A", "B♭", "B", "C", "D♭", "D", "E♭", "E", "F", "G♭", "G"][ord(pitch)]


func `$`*(sp: SpelledPitch): string =
  case sp.alteration
  of naDoubleFlat: return $sp.note & "𝄫"
  of naFlat: return $sp.note & "♭"
  of naNatural: return $sp.note
  of naSharp: return $sp.note & "#"
  of naDoubleSharp: return $sp.note & "𝄪"


func `$`*(sd: ScaleDegree): string =
  case sd.alteration
  of naDoubleFlat: return "𝄫" & $sd.number
  of naFlat: return "♭" & $sd.number
  of naNatural: return $sd.number
  of naSharp: return "#" & $sd.number
  of naDoubleSharp: return "𝄪" & $sd.number


func `$`*(k: Key): string =
  result = $k.tonic
  case k.mode
  of Ionian: discard
  of Dorian: result &= " Dorian"
  of Phrygian: result &= " Phrygian"
  of Lydian: result &= " Lydian"
  of Mixolydian: result &= " Mixolydian"
  of Aeolian: result &= " Minor"
  of Locrian: result &= " Locrian"


func chromaticDistanceFromTonic*(degree: ScaleDegree): int =
  ## Return the number of semitones from the tonic to the given scale degree.
  ## This ignores modes and key signatures and speaks in terms of intervals
  ## instead.
  if degree.number < 1 or degree.number > 7:
    raise newException(ValueError, "Invalid scale degree: " & $degree.number)

  result = MajorIntervals[(degree.number - 1)] + ord(degree.alteration)


func toPitch*(n: Note): Pitch =
  ## Return the chromatic pitch for the given natural note.
  cast[Pitch]((4 + MajorIntervals[(ord(n) + 5) mod 7]) mod 12)


func toPitch*(sp: SpelledPitch): Pitch =
  ## Return the chromatic pitch for a given spelled note.
  cast[Pitch]((ord(sp.note.toPitch) + ord(sp.alteration) + 12) mod 12)


func toNote*(p: Pitch): Note =
  ## Return the natural note for a given chromatic pitch.
  case p
  of Af, A: Note.A
  of Bf, B: Note.B
  of C: Note.C
  of Df, D: Note.D
  of Ef, E: Note.E
  of F: Note.F
  of Gf, G: Note.D


func toSpelledPitch*(p: Pitch): SpelledPitch =
  ## Get a SpelledPitch version of a chromatic Pitch. Note that this always
  ## uses the simplest natural or flat spelling for the Pitch (Pitch.Bf is
  ## always spelled as B♭, never C𝄫)
  SpelledPitch(
    note: p.toNote,
    alteration:
      case p
      of Af, Bf, Df, Ef, Gf: naFlat
      else: naNatural)


func parseMode*(str: string): Mode =
  case str.toLower.strip
  of "ionian", "major", "": Ionian
  of "dorian": Dorian
  of "phrygian": Phrygian
  of "lydian": Lydian
  of "mixolydian": Mixolydian
  of "aeolian", "minor": Aeolian
  of "locrian": Locrian
  else: raise newException(ValueError, "Unrecognized mode/scale: " & str)


func parseSpelledPitch*(str: string): SpelledPitch =
  try: result.note = parseEnum[Note]($str[0])
  except:
    raise newException(ValueError, str[0] & " is not a recognized pitch name.")

  result.alteration =
    case str[1..^1]
    of "𝄫", "bb", "♭♭": naDoubleFlat
    of "b", "♭", "f", "es": naFlat
    of "", "♮": naNatural
    of "s", "is", "#": naSharp
    of "𝄪", "##": naDoubleSharp
    else: raise newException(ValueError,
      str[1..^1] & " is not a recognized accidental.")


func parseScaleDegree*(str: string): ScaleDegree =
  try: result.number = parseInt($str[0])
  except:
    raise newException(ValueError, str[0] & " is not a valid scale degree.")

  result.alteration =
    case str[0..^2]
    of "𝄫", "bb", "♭♭": naDoubleFlat
    of "b", "♭", "f", "es": naFlat
    of "", "♮": naNatural
    of "s", "is", "#": naSharp
    of "𝄪", "##": naDoubleSharp
    else: raise newException(ValueError,
      str[0..^2] & " is not a recognized accidental.")


func ionianPitch*(key: Key, degreeNumber: int): Pitch =
  cast[Pitch]((ord(key.tonic.toPitch) + MajorIntervals[degreeNumber - 1]) mod 12)


func spellPitch*(k: Key, sd: ScaleDegree): SpelledPitch =
  ## Given a key and scale degree, spell it correctly in that key. For example,
  ## the ♭7 in C major is spelled B♭, the ♭7 in key of Db Locrian is spelled
  ## B𝄫, and the ♭7 in F# major is E.
  result.note = k.tonic.note + (sd.number - 1)
  let resultingPitch = ord(ionianPitch(k, sd.number)) + ord(sd.alteration)
  result.alteration = cast[NoteAlteration](
    resultingPitch - ord(result.note.toPitch))
  #[
  debugEcho "Spelling " & $sd & " in the key of " & $k & ":\n" &
    "\tsd.alteration: " & $ord(sd.alteration) &
    "\tkey.tonic.note: " & $k.tonic.note &
    "\tsd.number - 1: " & $(sd.number - 1) &
    "\tionianPitch: " & $ionianPitch(k, sd.number) &
    "\tord(ionianPitch): " & $ord(ionianPitch(k, sd.number)) &
    "\talteration: " & $ord(sd.alteration) &
    "\tresultingPitch: " & $resultingPitch &
    "\tresult.note: " & $(result.note) &
    "\tresult.alteration: " & $ord(result.alteration)
  ]#


func toScaleDegree*(k: Key, sp: SpelledPitch): ScaleDegree =
  ## Determine the ScaleDegree of a pitch according to how it is spelled and
  ## the key it is in. For example, Pitch.B will be the ♮2 in the key of A
  ## major, the ♭2 in the key of A# major, or the #1 in the key of B♭ major.
  result.number = sp.note - k.tonic.note + 1
  var distance = ord(sp.toPitch) - ord(ionianPitch(k, result.number))
  if distance < -2: distance += 12
  elif distance > 2: distance -= 12
  result.alteration = cast[NoteAlteration](distance)
  #[
  debugEcho "toScaleDegree: " & $sp & " in the key of " & $k & ":\n" &
    "\tsp.note: " & $sp.note &
    "\tkey.tonic.note: " & $k.tonic.note &
    "\tresult.number: " & $result.number &
    "\tsp.toPitch: " & $sp.toPitch &
    "\tionianPitch: " & $ionianPitch(k, result.number) &
    "\tord(sp.toPitch): " & $ord(sp.toPitch) &
    "\tord(ionianPitch): " & $ord(ionianPitch(k, result.number)) &
    "\tresult.alteration: " & $ord(result.alteration)
  ]#



func transpose*(k: Key, steps: int): Key =
  Key(
    tonic: toSpelledPitch(k.tonic.toPitch + steps),
    mode: k.mode)


when isMainModule:
  assert A + 1 == Bf
  assert A + 12 == A
  assert A - 1 == Af
  assert A + 14 == B

  assert C - A == 3
  assert A - G == 2
  assert G - A == 10

  assert Note.D - Note.B == 2
  assert Note.A - Note.C == 5
  assert Note.C - Note.A == 2

  assert chromaticDistanceFromTonic(ScaleDegree(number: 1, alteration: naNatural)) == 0
  assert chromaticDistanceFromTonic(ScaleDegree(number: 5, alteration: naNatural)) == 7
  assert chromaticDistanceFromTonic(ScaleDegree(number: 5, alteration: naDoubleFlat)) == 5
  assert chromaticDistanceFromTonic(ScaleDegree(number: 7, alteration: naFlat)) == 10
  assert chromaticDistanceFromTonic(ScaleDegree(number: 3, alteration: naDoubleSharp)) == 6

  assert $SpelledPitch(note: Note.B, alteration: naFlat) == "B♭"

  assert A == parseSpelledPitch("A").toPitch
  assert B == parseSpelledPitch("B").toPitch
  assert C == parseSpelledPitch("C").toPitch
  assert D == parseSpelledPitch("D").toPitch
  assert E == parseSpelledPitch("E").toPitch
  assert F == parseSpelledPitch("F").toPitch
  assert G == parseSpelledPitch("G").toPitch

  assert Af == parseSpelledPitch("Ab").toPitch
  assert G == parseSpelledPitch("Abb").toPitch
  assert B == parseSpelledPitch("Cb").toPitch
  assert Df == parseSpelledPitch("B##").toPitch

  assert "3" == $toScaleDegree(
    Key(tonic: parseSpelledPitch("B"), mode: Ionian),
    parseSpelledPitch("D#"))

  assert "1" == $toScaleDegree(
    Key(tonic: parseSpelledPitch("Gb"), mode: Ionian),
    parseSpelledPitch("Gb"))

  assert parseSpelledPitch("Cb").toPitch == parseSpelledPitch("B").toPitch

  assert "♭6" == $toScaleDegree(
    Key(tonic: parseSpelledPitch("Eb"), mode: Ionian),
    parseSpelledPitch("Cb"))

  assert "#5" == $toScaleDegree(
    Key(tonic: parseSpelledPitch("Eb"), mode: Ionian),
    parseSpelledPitch("B"))