import re
import argparse
import itertools
from enum import Enum

rule_jokers_enabled = False

class CardType(Enum):
    HighCard = 1
    OnePair = 2
    TwoPair = 3
    ThreeOfAKind = 4
    FullHouse = 5
    FourOfAKind = 6
    FiveOfAKind = 7

class Hand:
    def __init__(self,cards:str,bid:int) -> None:
        self.cards = cards
        self.bid = int(bid)
        if rule_jokers_enabled:
            self.type = self._find_type_joker(cards)
        else:
            self.type = self._find_type(cards)

    def _find_type(self,cards:str) -> CardType:
        # group cards based on card counts
        card_list = [*cards]
        card_list.sort()
        grouping = itertools.groupby(card_list,lambda x:x)
        lengths = [len(list(x[1])) for x in grouping]
        if 5 in lengths:
            return CardType.FiveOfAKind
        if 4 in lengths:
            return CardType.FourOfAKind
        if 3 in lengths and 2 in lengths:
            return CardType.FullHouse
        if 3 in lengths:
            return CardType.ThreeOfAKind
        if len([x for x in lengths if x == 2]) == 2:
            return CardType.TwoPair
        if 2 in lengths:
            return CardType.OnePair
        return CardType.HighCard
    
    def _find_type_joker(self,cards:str) -> CardType:
        try:
            joker_i = cards.index("J") 
        except ValueError:
            return self._find_type(cards)
        
        current_value = CardType.HighCard
        for new_card in [*(valid_card_list())]:
            if new_card == "J":
                continue
            test_cards = list(cards)
            test_cards[joker_i] = new_card
            new_value = self._find_type_joker("".join(test_cards))
            if new_value.value > current_value.value:
                current_value = new_value
        
        return current_value

    
    def sort_string(self):
        v = str(self.type.value) + ":" + "".join(["abcdefghijklmnoZ"[card_value(x)] for x in [*self.cards]])
        return v
    
    def __repr__(self) -> str:
        return f""



def valid_card_list() -> str:
    if rule_jokers_enabled:
        return "J23456789TQKA"
    return "23456789TJQKA"

def card_value(char:chr):
    return valid_card_list().index(char)

def main(line_list: list):
    hand_list = list()
    for l in line_list:
        card,bid = re.split(' +',l)
        hand = Hand(card,bid)
        hand_list.append(hand)
        #print(hand.sort_string())
    
    hand_list.sort(key=lambda x: x.sort_string())
    print(hand_list)

    rank_total = 0
    rank = 1
    for single_hand in hand_list:
        rank_total += rank * single_hand.bid
        rank += 1
    
    print(f"total {rank_total}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="day 1 solver")
    parser.add_argument("-input",type=str)
    parser.add_argument("-part",type=int)
    args = parser.parse_args()

    if args.part == 2:
        rule_jokers_enabled = True

    filename = args.input
    if filename == None:
        parser.print_help()
        exit(1)
    file = open(filename,'r')
    main([line.rstrip('\n') for line in file.readlines()])
    file.close()

it doesn't make any sense to pick different cards for the jokers, and that it's always the highest score to assign all jokers to the most frequent card.

inductive Card
  | two
  | three
  | four
  | five
  | six
  | seven
  | eight
  | nine
  | ten
  | jack
  | queen
  | king
  | ace
  deriving Repr, Ord, BEq

inductive Hand
  | mk : Card → Card → Card → Card → Card → Hand
  deriving Repr, BEq

private inductive Score
  | highCard
  | onePair
  | twoPair
  | threeOfAKind
  | fullHouse
  | fourOfAKind
  | fiveOfAKind
  deriving Repr, Ord, BEq

-- we need countCards in part 2 again, but there it has different types
private class CardList (η : Type) (χ : outParam Type) where
  cardList : η → List χ

-- similarly, we can implement Ord in terms of CardList and Score
private class Scorable (η : Type) where
  score : η → Score

private instance : CardList Hand Card where
  cardList := λ
    | .mk a b c d e => [a,b,c,d,e]

private def countCards {η χ : Type} (input :η) [CardList η χ] [Ord χ] [BEq χ] : List (Nat × χ) :=
  let ordered := (CardList.cardList input).quicksort
  let helper := λ (a : List (Nat × χ)) (c : χ) ↦ match a with
  | [] => [(1, c)]
  | a :: as =>
    if a.snd == c then
      (a.fst + 1, c) :: as
    else
      (1, c) :: a :: as
  List.quicksortBy (·.fst > ·.fst) $ ordered.foldl helper []

private def evaluateCountedCards : (l : List (Nat × α)) → Score
  | [_] => Score.fiveOfAKind -- only one entry means all cards are equal
  | (4,_) :: _ => Score.fourOfAKind
  | [(3,_), (2,_)] => Score.fullHouse
  | (3,_) :: _ => Score.threeOfAKind
  | [(2,_), (2,_), _] => Score.twoPair
  | (2,_) :: _ => Score.onePair
  | _ => Score.highCard

private def Hand.score (hand : Hand) : Score :=
  evaluateCountedCards $ countCards hand

private instance : Scorable Hand where
  score := Hand.score

instance {σ χ : Type} [Scorable σ] [CardList σ χ] [Ord χ] : Ord σ where
  compare (a b : σ) :=
    let comparedScores := Ord.compare (Scorable.score a) (Scorable.score b)
    if comparedScores != Ordering.eq then
      comparedScores
    else
      Ord.compare (CardList.cardList a) (CardList.cardList b)

private def Card.fromChar? : Char → Option Card
| '2' => some Card.two
| '3' => some Card.three
| '4' => some Card.four
| '5' => some Card.five
| '6' => some Card.six
| '7' => some Card.seven
| '8' => some Card.eight
| '9' => some Card.nine
| 'T' => some Card.ten
| 'J' => some Card.jack
| 'Q' => some Card.queen
| 'K' => some Card.king
| 'A' => some Card.ace
| _ => none

private def Hand.fromString? (input : String) : Option Hand :=
  match input.toList.mapM Card.fromChar? with
  | some [a, b, c, d, e] => Hand.mk a b c d e
  | _ => none

abbrev Bet := Nat

structure Player where
  hand : Hand
  bet : Bet
  deriving Repr

def parse (input : String) : Except String (List Player) := do
  let lines := input.splitOn "\n" |> List.map String.trim |> List.filter String.notEmpty
  let parseLine := λ (line : String) ↦
    if let [hand, bid] := line.split Char.isWhitespace |> List.map String.trim |> List.filter String.notEmpty then
      Option.zip (Hand.fromString? hand) (String.toNat? bid)
      |> Option.map (uncurry Player.mk)
      |> Option.toExcept s!"Line could not be parsed: {line}"
    else
      throw s!"Failed to parse. Line did not separate into hand and bid properly: {line}"
  lines.mapM parseLine

def part1 (players : List Player) : Nat :=
  players.quicksortBy (λ p q ↦ p.hand < q.hand)
  |> List.enumFrom 1
  |> List.foldl (λ r p ↦ p.fst * p.snd.bet + r) 0


------------------------------------------------------------------------------------------------------
-- Again a riddle where part 2 needs different data representation, why are you doing this to me? Why?
-- (Though, strictly speaking, I could just add "joker" to the list of cards in part 1 and treat it special)

private inductive Card2
  | joker
  | two
  | three
  | four
  | five
  | six
  | seven
  | eight
  | nine
  | ten
  | queen
  | king
  | ace
  deriving Repr, Ord, BEq

private def Card.toCard2 : Card → Card2
  | .two => Card2.two
  | .three => Card2.three
  | .four => Card2.four
  | .five => Card2.five
  | .six => Card2.six
  | .seven => Card2.seven
  | .eight => Card2.eight
  | .nine => Card2.nine
  | .ten => Card2.ten
  | .jack => Card2.joker
  | .queen => Card2.queen
  | .king => Card2.king
  | .ace => Card2.ace

private inductive Hand2
  | mk : Card2 → Card2 → Card2 → Card2 → Card2 → Hand2
  deriving Repr

private def Hand.toHand2 : Hand → Hand2
  | Hand.mk a b c d e => Hand2.mk a.toCard2 b.toCard2 c.toCard2 d.toCard2 e.toCard2

instance : CardList Hand2 Card2 where
  cardList := λ
    | .mk a b c d e => [a,b,c,d,e]

private def Hand2.score (hand : Hand2) : Score :=
  -- I could be dumb here and just let jokers be any other card, but that would be really wasteful
  -- Also, I'm pretty sure there is no combination that would benefit from jokers being mapped to
  -- different cards.
  -- and, even more important, I think we can always map jokers to the most frequent card and are
  -- still correct.
  let counted := countCards hand
  let (jokers, others) := counted.partition λ e ↦ e.snd == Card2.joker
  let jokersReplaced := match jokers, others with
  | (jokers, _) :: _ , (a, ac) :: as => (a+jokers, ac) :: as
  | _ :: _, [] => jokers
  | [], others => others
  evaluateCountedCards jokersReplaced

private instance : Scorable Hand2 where
  score := Hand2.score

private structure Player2 where
  bet : Bet
  hand2 : Hand2

def part2 (players : List Player) : Nat :=
  let players := players.map λ p ↦
    {bet := p.bet, hand2 := p.hand.toHand2 : Player2}
  players.quicksortBy (λ p q ↦ p.hand2 < q.hand2)
  |> List.enumFrom 1
  |> List.foldl (λ r p ↦ p.fst * p.snd.bet + r) 0

input = File.read("input.txt").lines
rank = {
	'J' => 1,
	'2' => 2,
	'3' => 3,
	'4' => 4,
	'5' => 5,
	'6' => 6,
	'7' => 7,
	'8' => 8,	
	'9' => 9,
	'T' => 10,
	# 'J' => 11,
	'Q' => 12,
	'K' => 13,
	'A' => 14
}

hand = input.map do |line|
	split = line.split
	weights = split[0].chars.map {|c| rank[c]}
	{weights, split[1].to_i}
end

hand.sort! do |a, b|
	a_rank = get_rank(a[0], true)
	b_rank = get_rank(b[0], true)
	
	# puts "#{a}-#{a_rank} #{b}-#{b_rank}"
	next  1 if a_rank > b_rank
	next -1 if b_rank > a_rank

	val = 0
	5.times do |i| 
		val =  1 if a[0][i] > b[0][i]
		val = -1 if b[0][i] > a[0][i]
		break unless val == 0
	end
	val
end

sum = 0
hand.each_with_index do |card, i|
	sum += card[1]*(i+1)
end
puts sum


def get_rank(card : Array(Int32), joker = false ) : Float64 | Int32
	aa = card.uniq

	if joker
		card = aa.map { |c|
			combo = card.map {|a| a == 1 ? c : a }
			{combo, get_rank(combo)}
		}.max_by {|a| a[1]}[0]
		aa = card.uniq
	end
	
	rank = 6 - aa.size
	case rank
	when 3
		return 3.5 if card.count(aa[0]) == 3
		return 3   if card.count(aa[0]) == 2
		return 3   if card.count(aa[1]) == 2
		return 3.5
	when 4
		return 4 if card.count(aa[0]) == 3 || card.count(aa[0]) == 2
		return 4.5
	else 
		return rank
	end
end

// Part 1
// ======

function part1(input) {
  const lines = input.replaceAll("\r", "").split("\n");
  const hands = lines.map((line) => line.split(" "));

  const sortedHands = hands.sort((a, b) => {
    const handA = calculateHandValue(a[0]);
    const handB = calculateHandValue(b[0]);

    if (handA > handB) {
      return -1;
    } else if (handA < handB) {
      return 1;
    } else {
      for (let i = 0; i < 5; i++) {
        const handACard = convertToNumber(a[0].split("")[i]);
        const handBCard = convertToNumber(b[0].split("")[i]);
        if (handACard > handBCard) {
          return 1;
        } else if (handACard < handBCard) {
          return -1;
        }
      }
    }
  });

  return sortedHands
    .filter((hand) => hand[0] != "")
    .reduce((acc, hand, i) => {
      return acc + hand[1] * (i + 1);
    }, 0);
}

function convertToNumber(card) {
  switch (card) {
    case "A":
      return 14;
    case "K":
      return 13;
    case "Q":
      return 12;
    case "J":
      return 11;
    case "T":
      return 10;
    default:
      return parseInt(card);
  }
}

function calculateHandValue(hand) {
  const dict = {};

  hand.split("").forEach((card) => {
    if (dict[card]) {
      dict[card] += 1;
    } else {
      dict[card] = 1;
    }
  });

  // 5
  if (Object.keys(dict).length === 1) {
    return 1;
  }

  // 4
  if (Object.keys(dict).filter((key) => dict[key] === 4).length === 1) {
    return 2;
  }

  // 3 + 2
  if (
    Object.keys(dict).filter((key) => dict[key] === 3).length === 1 &&
    Object.keys(dict).filter((key) => dict[key] === 2).length === 1
  ) {
    return 3;
  }

  // 3
  if (Object.keys(dict).filter((key) => dict[key] === 3).length === 1) {
    return 4;
  }

  // 2 + 2
  if (Object.keys(dict).filter((key) => dict[key] === 2).length === 2) {
    return 5;
  }

  // 2
  if (Object.keys(dict).filter((key) => dict[key] === 2).length === 1) {
    return 6;
  }

  return 7;
}

// Part 2
// ======

function part2(input) {
  const lines = input.replaceAll("\r", "").split("\n");
  const hands = lines.map((line) => line.split(" "));

  const sortedHands = hands.sort((a, b) => {
    const handA = calculateHandValuePart2(a[0]);
    const handB = calculateHandValuePart2(b[0]);

    if (handA > handB) {
      return -1;
    } else if (handA < handB) {
      return 1;
    } else {
      for (let i = 0; i < 5; i++) {
        const handACard = convertToNumberPart2(a[0].split("")[i]);
        const handBCard = convertToNumberPart2(b[0].split("")[i]);
        if (handACard > handBCard) {
          return 1;
        } else if (handACard < handBCard) {
          return -1;
        }
      }
    }
  });

  return sortedHands
    .filter((hand) => hand[0] != "")
    .reduce((acc, hand, i) => {
      console.log(acc, hand, i + 1);
      return acc + hand[1] * (i + 1);
    }, 0);
}

function convertToNumberPart2(card) {
  switch (card) {
    case "A":
      return 14;
    case "K":
      return 13;
    case "Q":
      return 12;
    case "J":
      return 1;
    case "T":
      return 10;
    default:
      return parseInt(card);
  }
}

function calculateHandValuePart2(hand) {
  const dict = {};

  let jokers = 0;

  hand.split("").forEach((card) => {
    if (card === "J") {
      jokers += 1;
      return;
    }
    if (dict[card]) {
      dict[card] += 1;
    } else {
      dict[card] = 1;
    }
  });

  // 5
  if (jokers === 5 || Object.keys(dict).length === 1) {
    return 1;
  }

  // 4
  if (
    jokers === 4 ||
    (jokers === 3 &&
      Object.keys(dict).filter((key) => dict[key] === 1).length >= 1) ||
    (jokers === 2 &&
      Object.keys(dict).filter((key) => dict[key] === 2).length === 1) ||
    (jokers === 1 &&
      Object.keys(dict).filter((key) => dict[key] === 3).length === 1) ||
    Object.keys(dict).filter((key) => dict[key] === 4).length === 1
  ) {
    return 2;
  }

  // 3 + 2
  if (
    (Object.keys(dict).filter((key) => dict[key] === 3).length === 1 &&
      Object.keys(dict).filter((key) => dict[key] === 2).length === 1) ||
    (Object.keys(dict).filter((key) => dict[key] === 2).length === 2 &&
      jokers === 1)
  ) {
    return 3;
  }

  // 3
  if (
    Object.keys(dict).filter((key) => dict[key] === 3).length === 1 ||
    (Object.keys(dict).filter((key) => dict[key] === 2).length === 1 &&
      jokers === 1) ||
    (Object.keys(dict).filter((key) => dict[key] === 1).length >= 1 &&
      jokers === 2) ||
    jokers === 3
  ) {
    return 4;
  }

  // 2 + 2
  if (
    Object.keys(dict).filter((key) => dict[key] === 2).length === 2 ||
    (Object.keys(dict).filter((key) => dict[key] === 2).length === 1 &&
      jokers === 1)
  ) {
    return 5;
  }

  // 2
  if (
    Object.keys(dict).filter((key) => dict[key] === 2).length === 1 ||
    jokers
  ) {
    return 6;
  }

  return 7;
}

export default { part1, part2 };