import re
from typing import List

import torch

log_file1 = "/data0/wangyin/workspace/FasterTransformer/tp1.log"
log_file2 = "/data0/wangyin/workspace/FasterTransformer/tp2.log"


class RowLog:
    def __init__(self, batch_idx: int, seq_idx: int, elements: List[float]) -> None:
        self.batch_idx = batch_idx
        self.seq_idx = seq_idx
        self.elements = elements

    def __str__(self) -> str:
        return f"batch_idx: {self.batch_idx}, seq_idx: {self.seq_idx}, elements: {self.elements}"

    def __repr__(self) -> str:
        return self.__str__()


class TensorLog:
    def __init__(
        self,
        name: str,
        seq_len: int,
        batch_size: int,
        dim: int,
        rows: List[RowLog] = [],
    ) -> None:
        self.name = name
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.dim = dim
        self.rows = rows

    def add_row(self, row: RowLog) -> None:
        self.rows.append(row)


def parse_log(file_path: str) -> List[TensorLog]:
    logs = []
    with open(file_path, "r") as f:
        lines = f.readlines().__iter__()

        try:
            next_line = lines.__next__()
            while next_line:
                match = re.search(r"^(.+)\[(\d+)\s(\d+)\s(\d+)\]$", next_line)
                if match:
                    current_log = TensorLog(
                        match.group(1),
                        int(match.group(2)),
                        int(match.group(3)),
                        int(match.group(4)),
                        [],
                    )
                    while True:
                        next_line = lines.__next__()
                        # checking the line only contains numbers and spaces
                        try:
                            elements = [
                                float(x.replace("b_", "").replace("s_", ""))
                                for x in next_line.split()
                            ]
                            current_log.add_row(
                                RowLog(int(elements[0]), int(elements[1]), elements[2:])
                            )
                        except ValueError:
                            break
                    # print(f'loaded log {current_log.name} with {len(current_log.rows)} rows')
                    logs.append(current_log)
                else:
                    next_line = lines.__next__()
        except StopIteration:
            pass
        print(f"totally loaded {len(logs)} logs")

    return logs


def compare_logs(log1: List[TensorLog], log2: List[TensorLog]) -> None:
    for i in range(min(len(log1), len(log2))):
        if log1[i].name != log2[i].name:
            print(f"log name not match: {log1[i].name} vs {log2[i].name}")
            return
        if log1[i].seq_len != log2[i].seq_len:
            print(f"log seq_len not match: {log1[i].seq_len} vs {log2[i].seq_len}")
            return
        if log1[i].batch_size != log2[i].batch_size:
            print(
                f"log batch_size not match: {log1[i].batch_size} vs {log2[i].batch_size}"
            )
            continue

        m1 = torch.tensor([row.elements for row in log1[i].rows])
        m2 = torch.tensor([row.elements for row in log2[i].rows])
        diff = torch.abs(m1 - m2)
        print(f"log [{log1[i].name}] diff: {diff}")


if __name__ == "__main__":
    logs1 = parse_log(log_file1)
    logs2 = parse_log(log_file2)
    compare_logs(logs1, logs2)
    print("done")
