from flatland.envs.agent_utils import RailAgentStatus
from flatland.envs.malfunction_generators import malfunction_from_params, MalfunctionParameters
from flatland.envs.observations import GlobalObsForRailEnv
from flatland.envs.rail_env import RailEnv
from flatland.envs.rail_generators import sparse_rail_generator
from flatland.envs.schedule_generators import sparse_schedule_generator
from flatland.utils.rendertools import RenderTool

import random
import sys
import os
import time
import msgpack
import json
from PIL import Image

import argparse as ap


def RandomTestParams(tid):
    seed = tid * 19997 + 997
    random.seed(seed)
    width = 50 + random.randint(0, 100)
    height = 50 + random.randint(0, 100)
    nr_cities = 4 + random.randint(0, (width + height) // 10)
    nr_trains = min(nr_cities * 20, 100 + random.randint(0, 100))
    max_rails_between_cities = 2
    max_rails_in_cities = 3 + random.randint(0, 5)
    malfunction_rate = 30 + random.randint(0, 100)
    malfunction_min_duration = 3 + random.randint(0, 7)
    malfunction_max_duration = 20 + random.randint(0, 80)
    return (
        seed, width, height, 
        nr_trains, nr_cities,
        max_rails_between_cities, max_rails_in_cities,
        malfunction_rate, malfunction_min_duration, malfunction_max_duration
        )


def RandomTestParams_small(tid):
    seed = tid * 19997 + 997
    random.seed(seed)

    nSize = random.randint(0,5)

    width = 20 + nSize * 5
    height = 20 + nSize * 5
    nr_cities = 2 + nSize // 2 + random.randint(0,2)
    nr_trains = min(nr_cities * 5, 5 + random.randint(0,5)) #, 10 + random.randint(0, 10))
    max_rails_between_cities = 2
    max_rails_in_cities = 3 + random.randint(0, nSize)
    malfunction_rate = 30 + random.randint(0, 100)
    malfunction_min_duration = 3 + random.randint(0, 7)
    malfunction_max_duration = 20 + random.randint(0, 80)
    return (
        seed, width, height, 
        nr_trains, nr_cities,
        max_rails_between_cities, max_rails_in_cities,
        malfunction_rate, malfunction_min_duration, malfunction_max_duration
        )



def ShouldRunTest(tid):
    return tid >= 7
    #return tid >= 3
    return True

                           

def create_test_env(fnParams, nTest, sDir):
    (seed, width, height,
    nr_trains, nr_cities, 
    max_rails_between_cities, max_rails_in_cities, 
    malfunction_rate, malfunction_min_duration, malfunction_max_duration) = fnParams(nTest)
    #if not ShouldRunTest(test_id):
    #    continue

    rail_generator = sparse_rail_generator(
        max_num_cities=nr_cities,
        seed=seed,
        grid_mode=False,
        max_rails_between_cities=max_rails_between_cities,
        max_rails_in_city=max_rails_in_cities,
    )


    

    #stochastic_data = {'malfunction_rate': malfunction_rate,
    #                    'min_duration': malfunction_min_duration,
    #                    'max_duration': malfunction_max_duration
    #                }

    stochastic_data = MalfunctionParameters(malfunction_rate=malfunction_rate,
                        min_duration=malfunction_min_duration,
                        max_duration=malfunction_max_duration
                    )





    observation_builder = GlobalObsForRailEnv()


    DEFAULT_SPEED_RATIO_MAP = {
        1.: 0.25,
        1. / 2.: 0.25,
        1. / 3.: 0.25,
        1. / 4.: 0.25}

    schedule_generator = sparse_schedule_generator(DEFAULT_SPEED_RATIO_MAP)

    for iAttempt in range(5):
        try:
            env = RailEnv(
                width=width,
                height=height,
                rail_generator=rail_generator,
                schedule_generator=schedule_generator,
                number_of_agents=nr_trains,
                malfunction_generator_and_process_data=malfunction_from_params(stochastic_data),
                obs_builder_object=observation_builder,
                remove_agents_at_target=True
                )
            obs = env.reset(random_seed = seed)
            break
        except ValueError as oErr:
            print("Error:", oErr)
            width += 5
            height += 5
            print("Try again with larger env: (w,h):", width, height)


    if not os.path.exists(sDir):
        os.makedirs(sDir)

    sfName = "{}/Level_{}.mpk".format(sDir, nTest)
    if os.path.exists(sfName):
        os.remove(sfName)
    env.save(sfName)

    sys.stdout.write(".")
    sys.stdout.flush()

    return env


#env = create_test_env(RandomTestParams_small, 0, "train-envs-small/Test_0")

def createEnvSet(nStart, nEnd, sDir, bSmall=True):

    #print("Generate small envs in train-envs-small:")
    print(f"Generate envs (small={bSmall}) in dir {sDir}:")

    sDirImages = "train-envs-small/images/"
    if not os.path.exists(sDirImages):
        os.makedirs(sDirImages)

    for test_id in range(nStart, nEnd, 1):
        env = create_test_env(RandomTestParams_small, test_id, sDir)

        oRender = RenderTool(env, gl="PILSVG")

        #oRender.env = env
        #oRender.set_new_rail()
        oRender.render_env()
        g2img = oRender.get_image()
        imgPIL = Image.fromarray(g2img)
        #imgPIL.show()

        imgPIL.save(sDirImages + "Level_{}.png".format(test_id))


    # print("Generate large envs in train-envs-1000:")

    # for test_id in range(100):
    #     create_test_env(RandomTestParams, test_id, "train-envs-1000/Test_0")



def merge(sfEpisode, sfEnv, sfEnvOut, bJson=False):
    if bJson:
        with open(sfEpisode, "rb") as fEp:
            oActions = json.load(fEp)
            oEp = {"actions":oActions}
            print("json oEp:", type(oEp), list(oEp.keys()))
    else:
        with open(sfEpisode, "rb") as fEp:
            oEp = msgpack.load(fEp)
            print("oEp:", type(oEp), list(oEp.keys()))
    
    with open(sfEnv, "rb") as fEnv:
        oEnv = msgpack.load(fEnv)
        print("oEnv:", type(oEnv), list(oEnv.keys()))
    
    # merge dicts
    oEnv2 = {**oEp, **oEnv}
    print("Merged keys:", list(oEnv2.keys()))
        
    with open(sfEnvOut, "wb") as fEnv:
        msgpack.dump(oEnv2, fEnv)
        
def printKeys1(sfEnv):
    with open(sfEnv, "rb") as fEnv:
        oEnv = msgpack.load(fEnv, encoding="utf-8")
        print(sfEnv, "keys:", list(oEnv.keys()))
        for sKey in oEnv.keys():
            print("key", sKey, len(oEnv[sKey]))
            if sKey == "shape":
                print("shape: ", oEnv[sKey] )


def printKeys(sfEnvs):
    try:
        for sfEnv in sfEnvs:
            printKeys1(sfEnv)
    except:
        # assume single env
        printKeys1(sfEnvs)

    



def main2():
    parser = ap.ArgumentParser(description='Generate envs, merge episodes into env files.')

    parser.add_argument("-c", '--createEnvs',  type=int, nargs=2, action="append",
        metavar=("nStart", "nEnd"),
        help='merge episode into env')

    parser.add_argument("-d", "--outDir", type=str, nargs=1, default="./test-envs-tmp")

    parser.add_argument("-m", '--merge',  type=str, nargs=3, action="append",
        metavar=("episode", "env", "output_env"),
        help='merge episode into env')
    
    parser.add_argument("-j", '--mergejson',  type=str, nargs=3, action="append",
        metavar=("json", "env", "output_env"),
        help='merge json actions into env, with key actions')
    
    parser.add_argument('-k', "--keys", type=str, action='append', nargs="+",
        help='print the keys in a file')

    args=parser.parse_args()
    print(args)

    if args.merge:
        print("merge:", args.merge)
        merge(*args.merge[0])

    if args.mergejson:
        print("merge json:", args.mergejson)
        merge(*args.mergejson[0], bJson=True)


    if args.keys:
        print("keys:", args.keys)
        printKeys(args.keys[0])

    if args.outDir:
        print("outDir", args.outDir)


    if args.createEnvs:
        print("create Envs - ", *args.createEnvs[0])
        createEnvSet(*args.createEnvs[0], sDir=args.outDir)
        

if __name__=="__main__":
    main2()

[((22, 4, 22, 21), 'random.seed', 'random.seed', ({(22, 16, 22, 20): 'seed'}, {}), '(seed)', False, 'import random\n'), ((42, 4, 42, 21), 'random.seed', 'random.seed', ({(42, 16, 42, 20): 'seed'}, {}), '(seed)', False, 'import random\n'), ((44, 12, 44, 31), 'random.randint', 'random.randint', ({(44, 27, 44, 28): '0', (44, 29, 44, 30): '5'}, {}), '(0, 5)', False, 'import random\n'), ((79, 21, 85, 5), 'flatland.envs.rail_generators.sparse_rail_generator', 'sparse_rail_generator', (), '', False, 'from flatland.envs.rail_generators import sparse_rail_generator\n'), ((95, 22, 98, 21), 'flatland.envs.malfunction_generators.MalfunctionParameters', 'MalfunctionParameters', (), '', False, 'from flatland.envs.malfunction_generators import malfunction_from_params, MalfunctionParameters\n'), ((104, 26, 104, 47), 'flatland.envs.observations.GlobalObsForRailEnv', 'GlobalObsForRailEnv', ({}, {}), '()', False, 'from flatland.envs.observations import GlobalObsForRailEnv\n'), ((113, 25, 113, 75), 'flatland.envs.schedule_generators.sparse_schedule_generator', 'sparse_schedule_generator', ({(113, 51, 113, 74): 'DEFAULT_SPEED_RATIO_MAP'}, {}), '(DEFAULT_SPEED_RATIO_MAP)', False, 'from flatland.envs.schedule_generators import sparse_schedule_generator\n'), ((140, 7, 140, 29), 'os.path.exists', 'os.path.exists', ({(140, 22, 140, 28): 'sfName'}, {}), '(sfName)', False, 'import os\n'), ((144, 4, 144, 25), 'sys.stdout.write', 'sys.stdout.write', ({(144, 21, 144, 24): '"""."""'}, {}), "('.')", False, 'import sys\n'), ((145, 4, 145, 22), 'sys.stdout.flush', 'sys.stdout.flush', ({}, {}), '()', False, 'import sys\n'), ((228, 13, 228, 91), 'argparse.ArgumentParser', 'ap.ArgumentParser', (), '', True, 'import argparse as ap\n'), ((23, 17, 23, 39), 'random.randint', 'random.randint', ({(23, 32, 23, 33): '(0)', (23, 35, 23, 38): '(100)'}, {}), '(0, 100)', False, 'import random\n'), ((24, 18, 24, 40), 'random.randint', 'random.randint', ({(24, 33, 24, 34): '(0)', (24, 36, 24, 39): '(100)'}, {}), '(0, 100)', False, 'import random\n'), ((25, 20, 25, 61), 'random.randint', 'random.randint', ({(25, 35, 25, 36): '(0)', (25, 38, 25, 60): '((width + height) // 10)'}, {}), '(0, (width + height) // 10)', False, 'import random\n'), ((28, 30, 28, 50), 'random.randint', 'random.randint', ({(28, 45, 28, 46): '(0)', (28, 48, 28, 49): '(5)'}, {}), '(0, 5)', False, 'import random\n'), ((29, 28, 29, 50), 'random.randint', 'random.randint', ({(29, 43, 29, 44): '(0)', (29, 46, 29, 49): '(100)'}, {}), '(0, 100)', False, 'import random\n'), ((30, 35, 30, 55), 'random.randint', 'random.randint', ({(30, 50, 30, 51): '(0)', (30, 53, 30, 54): '(7)'}, {}), '(0, 7)', False, 'import random\n'), ((31, 36, 31, 57), 'random.randint', 'random.randint', ({(31, 51, 31, 52): '(0)', (31, 54, 31, 56): '(80)'}, {}), '(0, 80)', False, 'import random\n'), ((48, 33, 48, 52), 'random.randint', 'random.randint', ({(48, 48, 48, 49): '(0)', (48, 50, 48, 51): '(2)'}, {}), '(0, 2)', False, 'import random\n'), ((51, 30, 51, 54), 'random.randint', 'random.randint', ({(51, 45, 51, 46): '(0)', (51, 48, 51, 53): 'nSize'}, {}), '(0, nSize)', False, 'import random\n'), ((52, 28, 52, 50), 'random.randint', 'random.randint', ({(52, 43, 52, 44): '(0)', (52, 46, 52, 49): '(100)'}, {}), '(0, 100)', False, 'import random\n'), ((53, 35, 53, 55), 'random.randint', 'random.randint', ({(53, 50, 53, 51): '(0)', (53, 53, 53, 54): '(7)'}, {}), '(0, 7)', False, 'import random\n'), ((54, 36, 54, 57), 'random.randint', 'random.randint', ({(54, 51, 54, 52): '(0)', (54, 54, 54, 56): '(80)'}, {}), '(0, 80)', False, 'import random\n'), ((136, 11, 136, 31), 'os.path.exists', 'os.path.exists', ({(136, 26, 136, 30): 'sDir'}, {}), '(sDir)', False, 'import os\n'), ((137, 8, 137, 25), 'os.makedirs', 'os.makedirs', ({(137, 20, 137, 24): 'sDir'}, {}), '(sDir)', False, 'import os\n'), ((141, 8, 141, 25), 'os.remove', 'os.remove', ({(141, 18, 141, 24): 'sfName'}, {}), '(sfName)', False, 'import os\n'), ((158, 11, 158, 37), 'os.path.exists', 'os.path.exists', ({(158, 26, 158, 36): 'sDirImages'}, {}), '(sDirImages)', False, 'import os\n'), ((159, 8, 159, 31), 'os.makedirs', 'os.makedirs', ({(159, 20, 159, 30): 'sDirImages'}, {}), '(sDirImages)', False, 'import os\n'), ((164, 18, 164, 46), 'flatland.utils.rendertools.RenderTool', 'RenderTool', (), '', False, 'from flatland.utils.rendertools import RenderTool\n'), ((170, 17, 170, 39), 'PIL.Image.fromarray', 'Image.fromarray', ({(170, 33, 170, 38): 'g2img'}, {}), '(g2img)', False, 'from PIL import Image\n'), ((195, 15, 195, 33), 'msgpack.load', 'msgpack.load', ({(195, 28, 195, 32): 'fEnv'}, {}), '(fEnv)', False, 'import msgpack\n'), ((203, 8, 203, 33), 'msgpack.dump', 'msgpack.dump', ({(203, 21, 203, 26): 'oEnv2', (203, 28, 203, 32): 'fEnv'}, {}), '(oEnv2, fEnv)', False, 'import msgpack\n'), ((207, 15, 207, 51), 'msgpack.load', 'msgpack.load', (), '', False, 'import msgpack\n'), ((26, 42, 26, 64), 'random.randint', 'random.randint', ({(26, 57, 26, 58): '0', (26, 60, 26, 63): '100'}, {}), '(0, 100)', False, 'import random\n'), ((49, 39, 49, 58), 'random.randint', 'random.randint', ({(49, 54, 49, 55): '0', (49, 56, 49, 57): '5'}, {}), '(0, 5)', False, 'import random\n'), ((186, 23, 186, 37), 'json.load', 'json.load', ({(186, 33, 186, 36): 'fEp'}, {}), '(fEp)', False, 'import json\n'), ((191, 18, 191, 35), 'msgpack.load', 'msgpack.load', ({(191, 31, 191, 34): 'fEp'}, {}), '(fEp)', False, 'import msgpack\n'), ((123, 55, 123, 95), 'flatland.envs.malfunction_generators.malfunction_from_params', 'malfunction_from_params', ({(123, 79, 123, 94): 'stochastic_data'}, {}), '(stochastic_data)', False, 'from flatland.envs.malfunction_generators import malfunction_from_params, MalfunctionParameters\n')]