import itertools
from paraview.simple import *
from paraview import servermanager


def set_legend_label(view_obj, series, legend_label):
    # ParaView uses this system where the data array keys and legend names
    # are stored togheter in the SeriesLabel array.
    # SeriesLabel is ['Temperature_T_[C] ( block=1)', 'Temp1', 'Current_Density_J,_[A_m^-2] ( block=1)', 'current_dens']
    new_series = []
    set_label = False
    for i in view_obj.SeriesLabel:
        if set_label:
            i = legend_label
            set_label = False
        else:
            if series in i:
                # We have to change the next item (label) in the array.
                set_label = True
        new_series.append(i)
    view_obj.SeriesLabel = new_series


def change_n_items(string_list_propery, new_array: list, start_index: int) -> list:
    """Replace a part of string_list_propery with new_array elements starting at start_index
    
    Args:
        string_list_propery: ParaView's StringListProperty
    """
    for new_item in new_array:
        string_list_propery[start_index] = str(new_item)
        start_index += 1
    return string_list_propery


def set_line_color(view_obj, series, rgb_255):
    # This is simmiliar to set_legend_label.
    # ['Current_Density_J,_[A_m^-2] ( block=1)', '0', '0', '0', ...]
    # Where the '0', '0', '0', are RGB values (from 0 to 1)
    rgb_1 = [j / 255 for j in rgb_255]
    set_label = False
    color_set = 0  # To keep track of whick RGB values have been set
    for i, view_obj_series_obj in enumerate(view_obj.SeriesColor):
        if series in view_obj_series_obj:
            # The current item is "Temperature", so replace 
            # the next next three items (RGB) in the array.
            new_series = change_n_items(view_obj.SeriesColor, rgb_1, i+1)
    # view_obj.SeriesColor = list(new_series)


def format_chart_view(chart_view):
    chart_view.ChartTitle = "Temperature probes"
    chart_view.ChartTitleBold = 1
    chart_view.LegendLocation = "TopLeft"
    chart_view.BottomAxisTitle = "Time, s"
    chart_view.LeftAxisTitle = "Temperature, °C"
    chart_view.LeftAxisTitleBold = 1


LINE_COLORS = [
    [59, 48, 58],  # Black
    [255, 63, 41],  # Red
    [0, 30, 255],  # Blue
    [12, 184, 89],  # Green
]

probe_count = 0
for i in range(len(GetSources())):
    obj = list(GetSources().items())[i][1]
    if type(obj).__name__ == "InterpolatedProbe":
        probe_count += 1

# Find the first 'QuartileChartView' in which the new point will be plotted
quartile_chart_view = None
active_layout = GetLayout()
all_views = GetViewsInLayout(active_layout)
for view in all_views:
    if type(view).__name__ == "QuartileChartView":
        quartile_chart_view = view
        break

workpiece = GetActiveSource()

es = ExtractSelection()
dataset = servermanager.Fetch(es)
# Sometimes the block index is not 0, dont know why.
# The solution is to loop over 20 indexes and just use the first non-empty one. 
for i in range(20):
    selection_block = dataset.GetBlock(i)
    if bool(selection_block):
        break
if not selection_block:
    raise RuntimeError("Block index could not be found for selection")

coord = selection_block.GetPoints().GetPoint(0)
# Also display the coordinates in the plot legend later
cood_legend = f" ({round(coord[0], 3)}, {round(coord[1], 3)}, {round(coord[2], 3)})"
Delete(es)

# InterpolatedProbe is a custom filter that combines "extract location" and "plot data over time"
cenos_probe = InterpolatedProbe(registrationName="cenos_probe", Input=workpiece)
cenos_probe.Location = list(coord)
cenos_probe.Mode = "Interpolate At Location"

plotDataOverTime1 = cenos_probe


if not quartile_chart_view:
    quartile_chart_view = CreateView("QuartileChartView")
    plotDataOverTime1Display = Show(
        plotDataOverTime1, quartile_chart_view, "QuartileChartRepresentation"
    )
    AssignViewToLayout(view=quartile_chart_view, layout=active_layout, hint=0)

    format_chart_view(quartile_chart_view)
else:
    # Plot already exists, lets just add the new data to it
    plotDataOverTime1Display = Show(
        plotDataOverTime1, quartile_chart_view, "QuartileChartRepresentation"
    )

# For some reason the labels have the block id in them aswell (Temperature_T_[C] (block=2))
# So we have to find the label by searching for one that contains "Temperature_T_[C]".
series_label = next(
    i for i in plotDataOverTime1Display.SeriesLabel if "Temperature_T_[C]" in i
)
plotDataOverTime1Display.SeriesVisibility = [series_label]

plotDataOverTime1Display.ShowRanges = 0
plotDataOverTime1Display.ShowQuartiles = 0

RenameSource("Probe " + str(probe_count + 1), plotDataOverTime1)
set_legend_label(
    plotDataOverTime1Display,
    "Temperature",
    "Probe " + str(probe_count + 1) + cood_legend,
)

# If the number of probes is greater than number of colors, just loop the colors
color = LINE_COLORS[probe_count % len(LINE_COLORS)]

# try:
# yeah I dont know why this doesnt work sometimes,
# just let the line be the default brown.
set_line_color(plotDataOverTime1Display, "Temperature", color)
# except:
#     pass