import random
import os
import time
import json
import subprocess
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

SOLUTION_EXT = "../submissions/accepted/jimmy_n3_ext"
SOLUTION = "../submissions/accepted/jimmy_n3"

CACHE_FILE = "status_cache.ndjson"

MAX_N = 200
MAX_M = 10**9
MAX_V = 10**9

case_num = 1

TMP_DATA_FILE = "tmpdata.tmp"

cache: Dict[int, Dict] = {}

@dataclass(frozen=True)
class Segment:
    start: int
    stop: int
    value: int

@dataclass(frozen=True)
class TestCase:
    n: int
    m: int
    x: int
    segments: Tuple[Segment]


def tmp_output(case: TestCase):
    with open(TMP_DATA_FILE, "w") as f:
        f.write(f"{case.m} {case.x} {case.n}\n")
        for segment in case.segments:
            f.write(f"{segment.start} {segment.stop} {segment.value}\n")


def output(description, case: TestCase):
    global case_num

    with open(f"secret/{case_num:02}-{description}.in", "wt") as f:
        f.write(f"{case.m} {case.x} {case.n}\n")
        for segment in case.segments:
            f.write(f"{segment.start} {segment.stop} {segment.value}\n")

    case_num += 1

def check_case(case: TestCase):
    if case.__hash__() in cache:
        return cache[case.__hash__()]

    tmp_output(case)

    with open(TMP_DATA_FILE) as f:
        res = subprocess.run(SOLUTION_EXT, capture_output=True, stdin=f)
        status = json.loads(res.stderr.strip())
        print(status)
        cache[case.__hash__()] = status
        return status

def dump_cache():
    with open(CACHE_FILE, "w") as f:
        for key, status in cache.items():
            v = {"key": key, "status": status}
            f.write(json.dumps(v) + "\n")       

def load_cache():
    with open(CACHE_FILE, "r") as f:
        for s in f.readlines():
            if s.strip():
                v = json.loads(s)
                cache[int(v["key"])] = v["status"]


def clean():
    for filename in os.listdir("secret"):
        if filename.endswith(".in") or filename.endswith(".ans"):
            os.remove(f"secret/{filename}")

def create_ans():
    slowest = 0.0

    for filename in os.listdir("secret"):
        if filename.endswith(".in"):
            ans_filename = filename.replace(".in", ".ans")
            start = time.time()
            os.system(f"{SOLUTION} <secret/{filename} >secret/{ans_filename}")
            elapsed = time.time() - start
            if elapsed > slowest:
                slowest = elapsed

    print(f"Slowest = {slowest:.2} s")

def read_manual(file_name):
    with open(file_name, "r") as f:
        m, x, n = map(int, f.readline().split(' '))
        segments = []
        for line in f.readlines():
            seg = tuple(map(int, line.split(' ')))
            segments.append(Segment(seg[0], seg[1], seg[2]))
    return TestCase(n, m, x, tuple(segments))

def random_segments(m, n, maxv) -> List[Segment]:
    x = set()
    while len(x) < 2*n:
        x.add(random.randint(0, m))
    x = list(sorted(x))
    return [Segment(x[i*2], x[i*2+1], random.randint(1, maxv)) for i in range(n)]

def generate_random_case(n=MAX_N, m=MAX_M, maxv=MAX_V, x=-1):
    if x < 0:
        x = random.randint(1, 2*m)

    upper_n = random.randint(0, n)
    lower_n = n - upper_n
    upper_segments = random_segments(m, upper_n, maxv)
    lower_segments = random_segments(m, lower_n, maxv)
    lower_segments = [Segment(seg.stop, seg.start, seg.value) for seg in lower_segments]
    
    segments = [*upper_segments, *lower_segments]
    assert len(segments) == n
    random.shuffle(segments)

    # print(m, x, segments[0].start, segments[-1].stop)
    return TestCase(n, m, x, tuple(segments))

def generate_random_case_similar(n=MAX_N, m=MAX_M, maxv=MAX_V, x=-1, var=5):
    # Generates cases where the endpoints on both sides are close to each other
    if x < 0:
        x = random.randint(1, 2*m)
    upper_n = n // 2 + random.randint(0, n // 2)
    lower_n = n - upper_n
    assert 0 < lower_n <= upper_n < n
    upper_segments = random_segments(m, upper_n, maxv)
    
    points = [point for seg in upper_segments for point in [seg.start, seg.stop]]

    ok = False
    while not ok:    
        lower_points = list(random.sample(points, lower_n*2))
        for i in range(len(lower_points)):
            lower_points[i] = min(max(0, lower_points[i] + random.randint(-var, var)), m)
        lower_points.sort()

        lower_segments = [Segment(lower_points[i*2+1], lower_points[i*2], random.randint(1, maxv)) for i in range(lower_n)]
        ok = not any(seg.start == seg.stop for seg in lower_segments)
        
    segments = [*upper_segments, *lower_segments]
    assert len(segments) == n
    random.shuffle(segments)

    return TestCase(n, m, x, tuple(segments))
    

def generate_complete_fill_case(n=MAX_N, m=MAX_M, maxv=MAX_V, x=-1):
    # Generates cases where the segments completely fill up the entire road on both sides
    if x < 0:
        x = random.randint(1, 2*m)    
    upper_n = n // 2 + random.randint(0, n // 2)
    lower_n = n - upper_n
    assert 0 < lower_n <= upper_n < n
    
    def get_points(n):
        points = [1+x for x in random.sample(range(m-1),n)]
        points.extend([0, m])
        points.sort()
        return points
        
    upper_points = get_points(upper_n)
    lower_points = get_points(lower_n)
    upper_segments = [Segment(upper_points[i], upper_points[i+1], random.randint(1, maxv)) for i in range(upper_n)]
    lower_segments = [Segment(lower_points[i+1], lower_points[i], random.randint(1, maxv)) for i in range(lower_n)]
    
    segments = [*upper_segments, *lower_segments]
    assert len(segments) == n
    random.shuffle(segments)

    return TestCase(n, m, x, tuple(segments))

def generate_random_small_cases(m, upper_n, lower_n, maxv=MAX_V, x=-1):
    if x < 0:
        x = random.randint(m// 2, 2*m - m//2)

    def get_points(n):
        points = [1+x for x in random.sample(range(m-1),n)]
        points.extend([0, m])
        points.sort()
        return points
        
    upper_points = get_points(upper_n)
    lower_points = get_points(lower_n)
    upper_segments = [Segment(upper_points[i], upper_points[i+1], random.randint(1, maxv)) for i in range(upper_n)]
    lower_segments = [Segment(lower_points[i+1], lower_points[i], random.randint(1, maxv)) for i in range(lower_n)]
    
    segments = [*upper_segments, *lower_segments]
    random.shuffle(segments)

    return TestCase(upper_n + lower_n, m, x, tuple(segments))



if __name__ == "__main__":
    load_cache()

    random.seed(0)
    clean()
    
    cases = []

    for manual_name in os.listdir("manual"):
        c = read_manual(f"manual/{manual_name}")
        status = check_case(c)
        cases.append((c, "manual"))

    cases.append((TestCase(0, 1000, 500, tuple([])), "empty"))
    cases.append((TestCase(2, MAX_M, MAX_M*2, tuple([Segment(0, MAX_M, MAX_V), Segment(MAX_M, 0, MAX_V)])), "max-answer"))
    cases.append((TestCase(2, MAX_M, MAX_M, tuple([Segment(0, 1, 5), Segment(MAX_M, MAX_M-1, 7)])), "endpoints-1"))
    cases.append((TestCase(2, MAX_M, MAX_M+1, tuple([Segment(0, 1, 5), Segment(MAX_M, MAX_M-1, 7)])), "endpoints-2"))
    cases.append((TestCase(3, 2000, 1000, tuple([Segment(100, 500, 90), Segment(900, 1500, 70), Segment(1800, 200, 1)])), "noturn"))

    # TODO: Add cases with a lot (only?) 1-length segments

    # RANDOM CASES TYPE 1

    wanted_bit_mask = set([1, 2, 4])
    wanted_flip_rev_mask = set([1, 2, 4, 8])    
    
    # Generate one large input that has a unique solution for each "solution characteristic"
    while wanted_bit_mask or wanted_flip_rev_mask:
        c = generate_random_case(m=random.randint(MAX_M/2, MAX_M))
        status = check_case(c)
        if status["unique"]:
            if status["case_bit_mask"] in wanted_bit_mask or status["flip_rev_mask"] in wanted_flip_rev_mask:
                wanted_bit_mask.discard(status["case_bit_mask"])
                wanted_flip_rev_mask.discard(status["flip_rev_mask"])
                print(f"{len(wanted_bit_mask)} {len(wanted_flip_rev_mask)} left")
                cases.append((c, "random"))
        

    # Generate some large cases, unique and non-unique, that overlaps between 70 and 90% of all segments
    num_large_segments_unique_left = 3
    num_large_segments_non_unique_left = 3
    while num_large_segments_unique_left or num_large_segments_non_unique_left:
        c = generate_random_case(m=random.randint(MAX_M/2, MAX_M))
        status = check_case(c)
        if status["segments"] >= MAX_N * 0.7 and status["segments"] < MAX_N * 0.9:
            if status["unique"] and num_large_segments_unique_left:
                num_large_segments_unique_left -= 1
                cases.append((c, "random"))
            if not status["unique"] and num_large_segments_non_unique_left:
                num_large_segments_non_unique_left -= 1
                cases.append((c, "random"))
    
        print(f"{num_large_segments_unique_left} {num_large_segments_non_unique_left} left")


    # RANDOM CASES TYPE 2 ("SIMILAR")

    num_large_segments_unique_left = 5
    num_large_segments_non_unique_left = 5
    while num_large_segments_unique_left or num_large_segments_non_unique_left:
        c = generate_random_case(m=random.randint(MAX_M/2, MAX_M), n = MAX_N-random.randint(0, 10))
        status = check_case(c)
        if status["unique"] and num_large_segments_unique_left:
            num_large_segments_unique_left -= 1
            cases.append((c, "similar"))
        if not status["unique"] and num_large_segments_non_unique_left:
            num_large_segments_non_unique_left -= 1
            cases.append((c, "similar"))
    
        print(f"{num_large_segments_unique_left} {num_large_segments_non_unique_left} left")

    # RANDOM CASES TYPE 3 ("FILLED")

    maxv = 10
    while maxv <= MAX_V:
        c = generate_random_case(m=random.randint(MAX_M/2, MAX_M), n = MAX_N-random.randint(0, 10))
        status = check_case(c)
        cases.append((c, "filled"))
        maxv *= 10

    # RANDOM CASE TYPE 4 (short segments, 1-2 in length)

    small_left = 5
    while small_left > 0:
        c = generate_random_small_cases(MAX_N // 2, MAX_N//2-random.randint(5, 30), MAX_N//2-random.randint(5, 30))
        status = check_case(c)
        cases.append((c, "small-segments"))
        small_left -= 1

    dump_cache()

    for (c, description) in cases:
        output(description, c)
    
    print("Running solutions...")
    create_ans()