import numpy as np

from civrealm.freeciv.players.diplomacy_actions import GOLD_SET

from .utils import expand_dim, noop, onehotifier_maker

resize = {
    "unit": 128,
    "city": 32,
    "others_unit": 128,
    "others_city": 64,
    "others_player": 10,
}

map_ops = {
    "status": onehotifier_maker(3),
    "terrain": onehotifier_maker(14),
    "extras": noop,
    "output": noop,
    "tile_owner": expand_dim(-1),
    "city_owner": expand_dim(-1),
    "unit": noop,
    "unit_owner": expand_dim(-1),
}

unit_ops = {
    "x": noop,
    "y": noop,
    "owner": noop,
    "type_attack_strength": noop,
    "type_defense_strength": noop,
    "type_firepower": noop,
    "type_build_cost": noop,
    "type_convert_time": noop,
    "type_obsoleted_by": onehotifier_maker(53),
    "type_hp": noop,
    "type_move_rate": noop,
    "type_vision_radius_sq": noop,
    "type_worker": noop,
    "type_can_transport": noop,
    "home_city": noop,
    "moves_left": noop,
    "health": noop,
    "veteran": noop,
    "upkeep_food": noop,
    "upkeep_shield": noop,
    "upkeep_gold": noop,
}

others_unit_ops = {
    "id": noop,
    "owner": noop,
    "x": noop,
    "y": noop,
    "type": onehotifier_maker(52),
    "veteran": noop,
    "transported": noop,
    "hp": noop,
    "activity": noop,
    "activity_tgt": noop,
    "transported_by": noop,
    "keep_activity": noop,
    "occupied": noop,
}

city_ops = {
    "id": noop,
    "x": noop,
    "y": noop,
    "owner": noop,
    "size": noop,
    "food_stock": noop,
    "granary_size": noop,
    "granary_turns": noop,
    "luxury": noop,
    "science": noop,
    "prod_food": noop,
    "surplus_food": noop,
    "prod_gold": noop,
    "surplus_gold": noop,
    "prod_shield": noop,
    "surplus_shield": noop,
    "prod_trade": noop,
    "surplus_trade": noop,
    "bulbs": noop,
    "city_waste": noop,
    "city_corruption": noop,
    "city_pollution": noop,
    "state": onehotifier_maker(5),
    "turns_to_prod_complete": noop,
    "prod_process": noop,
    "production_value": onehotifier_maker(120),
    "ppl_angry": noop,
    "ppl_unhappy": noop,
    "ppl_content": noop,
    "ppl_happy": noop,
    "city_radius_sq": noop,
    "buy_cost": noop,
    "shield_stock": noop,
    "before_change_shields": noop,
    "disbanded_shields": noop,
    "caravan_shields": noop,
    "last_turns_shield_surplus": noop,
    "improvements": noop,
}

others_city_ops = {
    "id": noop,
    "owner": noop,
    "size": noop,
    "style": onehotifier_maker(10),
    "capital": noop,
    "occupied": noop,
    "walls": noop,
    "happy": noop,
    "unhappy": noop,
    "improvements": noop,
}

others_player_ops = {
    "owner_id": noop,
    "love": onehotifier_maker(
        [
            "Genocidal",
            "Belligerent",
            "Hostile",
            "Uncooperative",
            "Uneasy",
            "Neutral",
            "Respectful",
            "Helpful",
            "Enthusiastic",
            "Admiring",
            "Worshipful",
            "-",
        ]
    ),
    "score": noop,
    "is_alive": onehotifier_maker(2),
    "diplomatic_state": onehotifier_maker(8),
    "techs": noop,
}
player_ops = {
    "turn": noop,
    "culture": noop,
    "reasearching_cost": noop,
    "gold": noop,
    "government": onehotifier_maker(6),
    "is_alive": noop,
    "luxury": noop,
    "net_income": lambda x: np.array([0], dtype=np.int32)
    if x is None
    else x * np.array([1], dtype=np.int32),
    "revolution_finishes": noop,
    "science": noop,
    "science_cost": noop,
    "score": noop,
    "target_government": onehotifier_maker(7),
    "tax": noop,
    "my_tech_goal": lambda x: onehotifier_maker(89)(88 if x == 253 else x),
    "tech_upkeep": noop,
    "techs_researched": noop,
    "total_bulbs_prod": noop,
    "turns_alive": noop,
    "no_humans": noop,
    "no_ais": noop,
    "research_progress": noop,
    "team_no": noop,
    "techs": noop,
}

rules_ops = {
}

dipl_ops = {
    "type": noop,
    "give_city": noop,
    "ask_city": noop,
    "give_gold": onehotifier_maker(len(GOLD_SET) + 1),
    "ask_gold": onehotifier_maker(len(GOLD_SET) + 1),
}

ops = {
    "rules": rules_ops,
    "player": player_ops,
    "map": map_ops,
    "city": city_ops,
    "unit": unit_ops,
    "others_unit": others_unit_ops,
    "others_city": others_city_ops,
    "others_player": others_player_ops,
    "dipl": dipl_ops,
}

obs_mutable_layout = {
    "city": {
        "id": (1,),
        "owner": (1,),
        "size": (1,),
        "x": (1,),
        "y": (1,),
        "food_stock": (1,),
        "granary_size": (1,),
        "granary_turns": (1,),
        "production_value": (120,),
        "city_radius_sq": (1,),
        "buy_cost": (1,),
        "shield_stock": (1,),
        "disbanded_shields": (1,),
        "caravan_shields": (1,),
        "last_turns_shield_surplus": (1,),
        "improvements": (68,),
        "luxury": (1,),
        "science": (1,),
        "prod_food": (1,),
        "surplus_food": (1,),
        "prod_gold": (1,),
        "surplus_gold": (1,),
        "prod_shield": (1,),
        "surplus_shield": (1,),
        "prod_trade": (1,),
        "surplus_trade": (1,),
        "bulbs": (1,),
        "city_waste": (1,),
        "city_corruption": (1,),
        "city_pollution": (1,),
        "state": (5,),
        "turns_to_prod_complete": (1,),
        "prod_process": (1,),
        "ppl_angry": (6,),
        "ppl_unhappy": (6,),
        "ppl_content": (6,),
        "ppl_happy": (6,),
        "before_change_shields": (1,),
    },
    "unit": {
        "owner": (1,),
        "health": (1,),
        "veteran": (1,),
        "x": (1,),
        "y": (1,),
        "type_rule_name": (52,),
        "type_attack_strength": (1,),
        "type_defense_strength": (1,),
        "type_firepower": (1,),
        "type_build_cost": (1,),
        "type_convert_time": (1,),
        "type_obsoleted_by": (53,),
        "type_hp": (1,),
        "type_move_rate": (1,),
        "type_vision_radius_sq": (1,),
        "type_worker": (1,),
        "type_can_transport": (1,),
        "home_city": (1,),
        "moves_left": (1,),
        "upkeep_food": (1,),
        "upkeep_shield": (1,),
        "upkeep_gold": (1,),
    },
    "others_city": {
        "id": (1,),
        "owner": (1,),
        "size": (1,),
        "improvements": (68,),
        "style": (10,),
        "capital": (1,),
        "occupied": (1,),
        "walls": (1,),
        "happy": (1,),
        "unhappy": (1,),
    },
    "others_unit": {
        "owner": (1,),
        "veteran": (1,),
        "x": (1,),
        "y": (1,),
        "id": (1,),
        "type": (52,),
        "occupied": (1,),
        "transported": (1,),
        "hp": (1,),
        "activity": (1,),
        "activity_tgt": (1,),
        "transported_by": (1,),
        "keep_activity": (1,),
    },
    "others_player": {
        "score": (1,),
        "is_alive": (2,),
        "love": (12,),
        "diplomatic_state": (8,),
        "techs": (87,),
    },
    "dipl": {
        "type": (20,),
        "give_city": (32,),
        "ask_city": (64,),
        "give_gold": (16,),
        "ask_gold": (16,),
    },
}


obs_immutalbe_layout = {
    "rules": {
        "build_cost": (120,),
    },
    "map": {
        "status": (84, 56, 3),
        "terrain": (84, 56, 14),
        "extras": (84, 56, 34),
        "output": (84, 56, 6),
        "tile_owner": (84, 56, 1),
        "city_owner": (84, 56, 1),
        "unit": (84, 56, 52),
        "unit_owner": (84, 56, 1),
    },
    "player": {
        "score": (1,),
        "is_alive": (1,),
        "turns_alive": (1,),
        "government": (6,),
        "target_government": (7,),
        "tax": (1,),
        "science": (1,),
        "luxury": (1,),
        "gold": (1,),
        "culture": (1,),
        "revolution_finishes": (1,),
        "science_cost": (1,),
        "tech_upkeep": (1,),
        "techs_researched": (1,),
        "total_bulbs_prod": (1,),
        "techs": (87,),
    },
}

action_layout = {
    "city": {
        "city_work_None_": 4,
        "city_unwork_None_": 4,
        "city_work_": 20,
        "city_unwork_": 20,
        "city_buy_production": 1,
        "city_change_specialist_": 3,
        "produce_Settlers": 1,
        "produce_Workers": 1,
        "produce_Engineers": 1,
        "produce_Warriors": 1,
        "produce_Phalanx": 1,
        "produce_Archers": 1,
        "produce_Legion": 1,
        "produce_Pikemen": 1,
        "produce_Musketeers": 1,
        "produce_Partisan": 1,
        "produce_Alpine": 1,
        "produce_Riflemen": 1,
        "produce_Marines": 1,
        "produce_Paratroopers": 1,
        "produce_Mech": 1,
        "produce_Horsemen": 1,
        "produce_Chariot": 1,
        "produce_Knights": 1,
        "produce_Dragoons": 1,
        "produce_Cavalry": 1,
        "produce_Armor": 1,
        "produce_Catapult": 1,
        "produce_Cannon": 1,
        "produce_Artillery": 1,
        "produce_Howitzer": 1,
        "produce_Fighter": 1,
        "produce_Bomber": 1,
        "produce_Helicopter": 1,
        "produce_Stealth": 2,
        "produce_Trireme": 1,
        "produce_Caravel": 1,
        "produce_Galleon": 1,
        "produce_Frigate": 1,
        "produce_Ironclad": 1,
        "produce_Destroyer": 1,
        "produce_Cruiser": 1,
        "produce_AEGIS": 1,
        "produce_Battleship": 1,
        "produce_Submarine": 1,
        "produce_Carrier": 1,
        "produce_Transport": 1,
        "produce_Cruise": 1,
        "produce_Nuclear": 2,
        "produce_Diplomat": 1,
        "produce_Spy": 1,
        "produce_Caravan": 1,
        "produce_Freight": 1,
        "produce_Explorer": 1,
        "produce_Leader": 1,
        "produce_Barbarian": 1,
        "produce_AWACS": 1,
        "produce_Airport": 1,
        "city_sell_improvement_Airport": 1,
        "produce_Aqueduct": 1,
        "city_sell_improvement_Aqueduct": 1,
        "produce_Bank": 1,
        "city_sell_improvement_Bank": 1,
        "produce_Barracks": 3,
        "city_sell_improvement_Barracks": 3,
        "produce_Cathedral": 1,
        "city_sell_improvement_Cathedral": 1,
        "produce_City": 1,
        "city_sell_improvement_City": 1,
        "produce_Coastal": 1,
        "city_sell_improvement_Coastal": 1,
        "produce_Colosseum": 1,
        "city_sell_improvement_Colosseum": 1,
        "produce_Courthouse": 1,
        "city_sell_improvement_Courthouse": 1,
        "produce_Factory": 1,
        "city_sell_improvement_Factory": 1,
        "produce_Granary": 1,
        "city_sell_improvement_Granary": 1,
        "produce_Harbor": 1,
        "city_sell_improvement_Harbor": 1,
        "produce_Hydro": 1,
        "city_sell_improvement_Hydro": 1,
        "produce_Library": 1,
        "city_sell_improvement_Library": 1,
        "produce_Marketplace": 1,
        "city_sell_improvement_Marketplace": 1,
        "produce_Mass": 1,
        "city_sell_improvement_Mass": 1,
        "produce_Mfg": 1,
        "city_sell_improvement_Mfg": 1,
        "city_sell_improvement_Nuclear": 1,
        "produce_Offshore": 1,
        "city_sell_improvement_Offshore": 1,
        "produce_Palace": 1,
        "produce_Police": 1,
        "city_sell_improvement_Police": 1,
        "produce_Port": 1,
        "city_sell_improvement_Port": 1,
        "produce_Power": 1,
        "city_sell_improvement_Power": 1,
        "produce_Recycling": 1,
        "city_sell_improvement_Recycling": 1,
        "produce_Research": 1,
        "city_sell_improvement_Research": 1,
        "produce_SAM": 1,
        "city_sell_improvement_SAM": 1,
        "produce_SDI": 1,
        "city_sell_improvement_SDI": 1,
        "produce_Sewer": 1,
        "city_sell_improvement_Sewer": 1,
        "produce_Solar": 1,
        "city_sell_improvement_Solar": 1,
        "produce_Space": 3,
        "produce_Stock": 1,
        "city_sell_improvement_Stock": 1,
        "produce_Super": 1,
        "city_sell_improvement_Super": 1,
        "produce_Supermarket": 1,
        "city_sell_improvement_Supermarket": 1,
        "produce_Temple": 1,
        "city_sell_improvement_Temple": 1,
        "produce_University": 1,
        "city_sell_improvement_University": 1,
        "produce_Apollo": 1,
        "produce_A": 1,
        "produce_Colossus": 1,
        "produce_Copernicus": 1,
        "produce_Cure": 1,
        "produce_Darwin": 1,
        "produce_Eiffel": 1,
        "produce_Great": 2,
        "produce_Hanging": 1,
        "produce_Hoover": 1,
        "produce_Isaac": 1,
        "produce_J": 1,
        "produce_King": 1,
        "produce_Leonardo": 1,
        "produce_Lighthouse": 1,
        "produce_Magellan": 1,
        "produce_Manhattan": 1,
        "produce_Marco": 1,
        "produce_Michelangelo": 1,
        "produce_Oracle": 1,
        "produce_Pyramids": 1,
        "produce_SETI": 1,
        "produce_Shakespeare": 1,
        "produce_Statue": 1,
        "produce_Sun": 1,
        "produce_United": 1,
        "produce_Women": 1,
        "produce_Coinage": 1,
    },
    "unit": {
        "transform_terrain": 1,
        "mine": 1,
        "cultivate": 1,
        "plant": 1,
        "fortress": 1,
        "airbase": 1,
        "irrigation": 1,
        "keep_activity": 1,
        "paradrop": 1,
        "build_city": 1,
        "join_city": 1,
        "fortify": 1,
        "build_road": 1,
        "build_railroad": 1,
        "pillage": 1,
        "set_homecity": 1,
        "airlift": 1,
        "upgrade": 1,
        "deboard": 1,
        "board": 1,
        "unit_unload": 1,
        "cancel_order": 1,
        "goto_": 8,
        "attack_": 8,
        "conquer_city_": 8,
        "spy_bribe_unit_": 8,
        "spy_steal_tech_": 8,
        "spy_sabotage_city_": 8,
        "hut_enter_": 8,
        "embark_": 8,
        "disembark_": 8,
        "trade_route_": 9,
        "marketplace_": 9,
        "embassy_stay_": 8,
        "investigate_spend_": 8,
    },
    "dipl": {
        "stop_negotiation_": 1,
        "accept_treaty_": 1,
        "cancel_treaty_": 1,
        "cancel_vision_": 1,
        "add_clause_ShareMap_": 2,
        "remove_clause_ShareMap_": 2,
        "add_clause_ShareSeaMap_": 2,
        "remove_clause_ShareSeaMap_": 2,
        "add_clause_Vision_": 2,
        "remove_clause_Vision_": 2,
        "add_clause_Embassy_": 2,
        "remove_clause_Embassy_": 2,
        "add_clause_Ceasefire_": 2,
        "remove_clause_Ceasefire_": 2,
        "add_clause_Peace_": 2,
        "remove_clause_Peace_": 2,
        "add_clause_Alliance_": 2,
        "remove_clause_Alliance_": 2,
        "trade_tech_clause_Advance_": 174,
        "remove_clause_Advance_": 174,
        "trade_gold_clause_TradeGold_": len(GOLD_SET) * 2,
        "remove_clause_TradeGold_": len(GOLD_SET) * 2,
        "trunc_trade_city_clause_TradeCity_": resize["city"] + resize["others_city"],
        "trunc_remove_clause_TradeCity_": resize["city"] + resize["others_city"],
    },
    "gov": {
        "change_gov_Anarchy": 1,
        "change_gov_Despotism": 1,
        "change_gov_Monarchy": 1,
        "change_gov_Communism": 1,
        "change_gov_Republic": 1,
        "change_gov_Democracy": 1,
        "set_sci_luax_tax": 66,
    },
    "tech": {
        "research": 87,
    },
}

actor_type_list = ["city", "unit", "gov", "tech", "dipl", "turn done"]
mutables_others_have = ["unit", "city"]
immutables_others_have = ["player", "tech"]
mutable_fields = mutables_others_have + [
    "others" + field for field in mutables_others_have + immutables_others_have
]

immutables_only_i_have = ["map", "rules", "gov"]
immutable_fields = immutables_others_have + immutables_only_i_have

default_tensor_config = {
    "filter_observation": [
        "map",
        "unit",
        "player",
        "rules",
        "city",
        "dipl",
    ],
    "resize": resize,
    "obs_immutable_layout": obs_immutalbe_layout,
    "obs_mutable_layout": obs_mutable_layout,
    "obs_ops": ops,
    "actor_type_list": actor_type_list,
    "action_layout": action_layout,
    "mutables_others_have": mutables_others_have,
    "immutables_others_have": immutables_others_have,
    "mutable_fields": mutable_fields,
    "immutable_fields": immutable_fields,
}
