import os
import json
from pprint import pprint
import paraview.simple as pw  # type: ignore
from typing import Tuple, Dict, List

CONFIG_PATH = os.getenv("CENOS_CONFIG")
CASE_PATH = os.getcwd()
generated_conf_path = os.path.join(CONFIG_PATH, "jsons/_generated")
vis_conf_path = os.path.join(generated_conf_path, "visualization.json")
results_json_path = os.path.join(CASE_PATH, "results/results.json")

RES_FILE = os.path.join(CASE_PATH, "results/resFile.0.case")
RESULTS_JSON = json.load(open(results_json_path, encoding="utf-8"))
VIZ_CONF = json.load(open(vis_conf_path, encoding="utf-8"))
LAST_PARAMS = json.load(open(os.path.join(CASE_PATH, "last-params.json")))

CENOS_APP = RESULTS_JSON["general"]["app"]
DIM = RESULTS_JSON["general"]["dimension"]
IS_AXISYMMETRIC = RESULTS_JSON["general"]["isAxisymmetric"]


def find_state_file():
    for file in os.listdir(CASE_PATH):
        if file.endswith(".pvsm"):
            return os.path.join(CASE_PATH, file)


def load_ensight_file(file: str, label="Simulation Data"):
    # Read data with EnSight Reader.
    simulation_data = pw.EnSightReader(CaseFileName=file)
    pw.RenameSource(label, simulation_data)

    animation_scene1 = pw.GetAnimationScene()
    animation_scene1.UpdateAnimationUsingDataTimeSteps()
    # This Show function sort of initializes the reader, without it,
    # the composite data wil be 0
    pw.Show(simulation_data, pw.GetActiveView())
    pw.Hide(simulation_data)
    return simulation_data


def goto_last_timestep() -> None:
    animation_scene1 = pw.GetAnimationScene()
    animation_scene1.GoToLast()


def get_block_index_map(simulation_data) -> Dict[str, int]:
    """Return a map of block names and their IDs in simulation_data

    Example:
    {
        'Workpiece Surface': 0,
        'Workpiece': 1,
        'Inductor': 2,
        'Air': 3
    }
    """
    composite_data = simulation_data.GetDataInformation().GetCompositeDataInformation()
    block_count = composite_data.GetNumberOfChildren()

    block_name_id_map = {}
    for i in range(block_count):
        name = composite_data.GetName(i)
        block_name_id_map[name] = i
    return block_name_id_map


def hex_color_to_0_1(hex: str) -> Tuple[float, float, float]:
    """Convert hex color to 0-1 range rgb tuple"""
    hex = hex.strip("#")
    c_0_1 = tuple(int(hex[i : i + 2], 16) / 255 for i in (0, 2, 4))
    return c_0_1


def format_fields(view) -> None:
    """Set field colormaps, colorbar title, units and formatting."""
    for post_val in RESULTS_JSON["fieldFormatting"]:
        field_key = post_val["key"]

        field_safename = post_val["safeName"]
        lut = pw.GetColorTransferFunction(field_safename)
        lut_colorbar = pw.GetScalarBar(lut, view)

        # If the field is real or imaginary, add the im and re text.
        component = ""
        if "symbol" in post_val:
            if r"{im}" in post_val["symbol"]:
                component = r"$_{im}$ "
            if r"{re}" in post_val["symbol"]:
                component = r"$_{re}$ "

        if post_val["units"] == "%":
            title = f'{post_val["name"]}{component}, {post_val["units"]}'
        else:
            title = f'{post_val["name"]}{component}, ${post_val["units"]}$'

        lut_colorbar.Title = title
        lut_colorbar.ComponentTitle = ""

        if field_key in VIZ_CONF["field_colormaps"]:
            colormap_obj = VIZ_CONF["field_colormaps"][field_key]

            if "colormap" in colormap_obj:
                lut.ApplyPreset(colormap_obj["colormap"], True)

            if "valueFormat" in colormap_obj:
                lut_colorbar.AutomaticLabelFormat = 0
                lut_colorbar.LabelFormat = colormap_obj["valueFormat"]
                lut_colorbar.RangeLabelFormat = colormap_obj["valueFormat"]
            # Without this, a blank colorbar is shown.
        pw.HideScalarBarIfNotNeeded(lut, view)


class ExtractBlock:
    """A simplified interface for ParaView's ExtractBlock objects

    Using simulation_data and block_map in the constructor,
    it abstracts away the BlockIndices interface.

    Note that this shit only extracts a block for one render view
    """

    def __init__(self, simulation_data, block_map, render_view):
        self.extract_block = pw.ExtractBlock(Input=simulation_data)
        self.block_map = block_map
        self.blocks = []
        self.role = None
        self.render_view = render_view
        self.extract_block_display = pw.Show(
            self.extract_block, render_view, "UnstructuredGridRepresentation"
        )
        # This just lightens up the block, so it isnt so dark on the white background.
        self.extract_block_display.Ambient = 0.05

    def show_in_view(self, view) -> None:
        """Show the block in another view. This function overwrites the view object."""
        self.extract_block_display = pw.Show(
            self.extract_block, view, "UnstructuredGridRepresentation"
        )
        # This just lightens up the block, so it isnt so dark on the white background.
        dp = pw.GetDisplayProperties(proxy=self.extract_block, view=self.render_view)
        if dp.ColorArrayName[0] != None:
            pw.ColorBy(self.extract_block_display, tuple(dp.ColorArrayName))
        else:
            pw.ColorBy(self.extract_block_display, None)

        self.extract_block_display.Ambient = dp.Ambient
        self.extract_block_display.AmbientColor = dp.AmbientColor
        self.extract_block_display.DiffuseColor = dp.DiffuseColor
        self.extract_block_display.Representation = dp.Representation
        # self.extract_block_display.SetRepresentationType(dp.Representation)
        self.render_view = view

    def add_block(self, name: str) -> None:
        self.__check_block_in_map(name)
        self.blocks.append(name)
        self.__set_block_indices()

    def has_field(self, field: str) -> bool:
        """Check if specified field actually exists in the block data"""
        simulation_point_data = self.extract_block.PointData.keys()
        # The fields that are in simulation data exists for all blocks, so check
        # if they all values are 0, and if so, assume it is not in the block.
        if field in simulation_point_data:
            value_range = self.extract_block.PointData[field].GetRange()
            if (value_range[0] == 0.0) and (value_range[1] == 0.0):
                return False
            else:
                return True
        else:
            # TODO: Maybe should log this or something
            return False

    def show_field(self, field: str) -> None:
        """Set a certian point field to be represented"""
        pw.ColorBy(self.extract_block_display, ("POINTS", field))
        self.extract_block_display.SetScalarBarVisibility(self.render_view, True)
        self.extract_block_display.RescaleTransferFunctionToDataRange(True, False)

    def set_color(self, hex_color: str) -> None:
        colors_0_1 = hex_color_to_0_1(hex_color)
        self.extract_block_display.AmbientColor = colors_0_1
        self.extract_block_display.DiffuseColor = colors_0_1

    def set_representation(self, repr: str) -> None:
        self.extract_block_display.SetRepresentationType(repr)

    def hide(self) -> None:
        pw.Hide(self.extract_block, self.render_view)

    @property
    def label(self) -> str:
        return self._label

    @label.setter
    def label(self, label: str) -> None:
        self._label = label
        pw.RenameSource(label, self.extract_block)

    def __check_block_in_map(self, name):
        if name not in self.block_map:
            raise RuntimeError(
                f"ExtractBlock set_blocks: {name} not found in block_map: {self.block_map}"
            )

    def __set_block_indices(self):
        """Sets all of the blocks in self.extract_block"""
        # ParaView treats "0" as all blocks
        block_indexes = [self.block_map[name] + 1 for name in self.blocks]
        self.extract_block.BlockIndices = block_indexes


def style_block(ext_block: ExtractBlock, VIZ_CONF: dict, role: str) -> None:
    assert role in VIZ_CONF["roleStyle"], f"Role '{role}' was not found in VIZ_CONF"

    role_config = VIZ_CONF["roleStyle"][role]

    ext_block.set_color(role_config["color"])

    # Show the most relevant (first) available field.
    relevant_fields = role_config["relevantFields"]
    for field in relevant_fields:
        if ext_block.has_field(field):
            ext_block.show_field(field)
            break
    # The block will show solid color by default if no field was set

    if role_config["showByDefault"] == False:
        ext_block.hide()


def extract_default_blocks(sim_data, view) -> List[ExtractBlock]:
    """Extract block from simulation_data using result json and visualization config

    Blocks are extracted only if they exists in the simulation_data AND
    are defined in visualization.yaml config.

    Returns a list of extracted blocks.
    """
    extracted_blocks = []

    ensight_block_map = get_block_index_map(sim_data)

    grouped_blocks = []
    for group_name in RESULTS_JSON["blockGroups"].keys():
        group_obj = RESULTS_JSON["blockGroups"][group_name]
        # Extract block only if that role is defined in roleStyle.
        if group_obj["role"] in VIZ_CONF["roleStyle"]:
            ext_block = ExtractBlock(sim_data, ensight_block_map, view)
            ext_block.label = group_name

            for block_name in group_obj["blocks"]:
                ext_block.add_block(block_name)
                grouped_blocks.append(block_name)
            style_block(ext_block, VIZ_CONF, group_obj["role"])
            extracted_blocks.append(ext_block)

    for json_block in RESULTS_JSON["blocks"]:
        role = json_block["role"]
        # Extract block only if that role is defined in roleStyle.
        if role in VIZ_CONF["roleStyle"]:
            label = json_block["label"]
            # attempt to extract block only if it is the resFile and
            # do not add any block that were already extracted in a group
            if (label in ensight_block_map) and (label not in grouped_blocks):
                ext_block = ExtractBlock(sim_data, ensight_block_map, view)
                ext_block.role = role
                ext_block.add_block(label)
                ext_block.label = label
                style_block(ext_block, VIZ_CONF, role)
                extracted_blocks.append(ext_block)

    return extracted_blocks


def revolve_blocks(blocks: List[ExtractBlock], old_view, view) -> List[ExtractBlock]:
    for block in blocks:
        dp = pw.GetDisplayProperties(proxy=block.extract_block, view=old_view)
        if dp.Visibility == True:
            # Only extract visible blocks
            simulation_data = block.extract_block.Input
            # Create3Dfrom2D is a custom CENOS filter located in paraivew/filters
            revolved_block = pw.Create3Dfrom2D(Input=simulation_data)
            revolved_block.BlockIndices = block.extract_block.BlockIndices
            revolved_block.Angle = 240.0
            revolved_block.Resolution = 100

            revolved_block_display = pw.Show(
                revolved_block, view, "GeometryRepresentation"
            )

            if dp.ColorArrayName[0] != None:
                pw.ColorBy(revolved_block_display, tuple(dp.ColorArrayName))
            else:
                pw.ColorBy(revolved_block_display, None)

            revolved_block_display.Ambient = dp.Ambient
            revolved_block_display.AmbientColor = dp.AmbientColor
            revolved_block_display.DiffuseColor = dp.DiffuseColor
            pw.RenameSource(block.label + " 3D", revolved_block)


def create_secondary_view():
    """Split the active layout vertically and create a new RenderView to the right

    Lighting is added automatically.
    """
    layout1 = pw.GetLayout()
    layout1.SplitHorizontal(0, 0.5)
    render_view_2 = pw.CreateView("RenderView")
    layout1.AssignView(2, render_view_2)
    add_lighting(render_view_2)
    format_fields(render_view_2)
    return render_view_2


def create_axisymmetric_view(blocks, old_view, new_view):
    revolve_blocks(blocks, old_view, new_view)
    # The deafult camera positton assumes the z axis as up,
    # but in this case the y axis must be up.
    new_view.CameraPosition = [0.143212, 0.0907398, 0.1053643]
    new_view.CameraFocalPoint = [-6.6504e-18, 0.003499951, -0.001671941]
    new_view.CameraViewUp = [-0.362247, 0.898567, -0.2476959]

    return new_view


def load_plugins() -> None:
    """Loads all plugins in the specified directory"""
    directory = os.path.join(CONFIG_PATH, "paraview\plugins")
    if os.path.isdir(directory):
        for file in os.listdir(directory):
            if not file.startswith("_"):
                pw.LoadPlugin(os.path.join(directory, file))


def gotoLastTimestep() -> None:
    animationScene1 = pw.GetAnimationScene()
    animationScene1.UpdateAnimationUsingDataTimeSteps()
    animationScene1.GoToLast()


def set_camera_position(view) -> None:
    if DIM == 3:
        view.InteractionMode = "3D"
        view.CameraFocalPoint = [0, 0, 0]
        view.CameraViewUp = [-0.1, 0, 1]
        view.CameraPosition = [490, -98, 204]
    if DIM == 2:
        view.InteractionMode = "2D"
    view.ResetCamera()


def add_case_name(view) -> None:
    case_name = os.path.basename(os.path.normpath(CASE_PATH))
    text1 = pw.Text()
    text1.Text = case_name
    text1Display = pw.Show(text1, view, "TextSourceRepresentation")
    pw.RenameSource("Case Name", text1)


def add_lighting(view):
    view.UseLight = 0  # Disable default light kit

    def new_light(active_view, intensity, radius, posx, posy, posz):
        light = pw.AddLight(view=active_view)
        light.Coords = "Camera"
        light.Position = [posx, posy, posz]
        light.Intensity = intensity
        light.Radius = radius

    light0 = new_light(view, 0.63, 50.0, -1, 0, 1)  # from front left 45
    light1 = new_light(view, 0.63, 50.0, 1, 0, 1)  # from front right 45
    light2 = new_light(view, 0.28, 60.0, 0, 1, 0)  # from top
    light3 = new_light(view, 0.28, 60.0, 0, -1, 0)  # from bottom
    light4 = new_light(view, 0.18, 10.0, 0, 0, -1)  # Backlight accent
    light1 = pw.GetLight(0, view)  # Hide "light tooklit" sidebar
    pw.Hide3DWidgets(proxy=light1)


def createAxialView():
    layout1 = pw.GetLayout()
    layout1.SplitHorizontal(0, 0.5)
    renderView2 = pw.CreateView("RenderView")
    layout1.AssignView(2, renderView2)
    return renderView2