import copy
import os
import math
import warnings
import re
import weakref
import multiprocessing
from ipywidgets import IntProgress

import numpy as np
from IPython.display import Javascript, display # Part of Hack #1
import plotly as plotly
import plotly.graph_objects as go
import pint

import pozo
import pozo.renderers as pzr
import pozo.themes as pzt
import pozo.units as pzu

re_space = re.compile(' ')
re_power = re.compile(r'\*\*')

def javascript():
    add_scroll = '''var css = '.plot-container { overflow: auto; }',
head = document.getElementsByTagName('head')[0],
style = document.createElement('style');

style.type = 'text/css';
if (style.styleSheet){
  style.styleSheet.cssText = css;
} else {
  style.appendChild(document.createTextNode(css));
}
console.log("Executed javascript()")
head.appendChild(style);'''
    display(Javascript(add_scroll))

javascript()

# TODO FOR CROSSPLOT MUST X Y AND ALL COLOR HAVE THE SAME SHAPE OR ELSE BREAK INTERACTIVITY

# You can create your own dictionary like 'defaults' and
# construct a Style() with it (pass it in as `template`)
# Specifying values that are marked as "generated" will
# lead to undefined behavior, probably be overwritten.
# It may be better to inherit the Style() class and
# define your own class that overrides the generators.
# Furthermore, certain keywords are marked as required,
# since the generators depend on them and expect them.

defaults = dict(
    # Non-Plotly style values, used to generate Plotly layout values. These are all "required".
    track_margin = 15,      # margin between each track
    depth_axis_width = 35,  # size of depth axis if it's not the first thing, will also have margin above
    axis_label_height = 80, # size of each xaxis when stacked
    y_annotation_gutter = 10, # space for writing text

    # Plotly style values, structured like a plotly layout dictionary.
    # Important to know that `xaxes_template` will be used to create several `xaxis1`,`xaxis2`, `xaxisN` keys.
    # `xaxis_template` will not make it into the final layout dictionary passed to Plotly.
    plotly = dict(
        height = 900,          # required
        showlegend = True,
        legend = dict(
            yanchor = "top",
            y = 1,
            xanchor = "left",
#           x = 0,             # generated
        ),
        margin = dict(l=15, r=15, t=5, b=5),
        plot_bgcolor = '#FFFFFF',
        paper_bgcolor = '#FFFFFF',
#       width=?                # generated

        yaxis1 = dict(
                visible=True, # Would hdie gridelines too, so hide ticks and labels instead of axis.
                showgrid=True,
                zeroline=False,
                showline=False,
                position=0,
                gridcolor="#f0f0f0",
#               domain=[?,?],  # generated
#               maxallowed=,   # generated
#               minallowed=,   # generated
#               range=[?,?],   # generated
        ),

        xaxis_template = dict(
            showgrid=False, # TODO this is tough, grid where?
            zeroline=False,
            gridcolor="#f0f0f0", # this needs to be generated
            showline=True,
            linewidth=2,
            tickmode= 'auto',
            nticks= 0,
            ticks="outside",
            tickwidth=1,
            ticklen=6,
            tickangle=55,
#           side=['']              # generated
#           tickcolor=='#???',     # generated
#           linecolor='#???',      # generated
#           title=dict(
#               text="???",        # generated
#               font=dict(
#                   color="#???"   # generated
#           ),
#           tickfont=dict(
#               color="#???"       # generated
#           )
#           domain = [?,?]         # generated
#           anchor=""              # generated
#           autoshift=             # Doesn't work for x-axes!
#           position=              # generated
#           overlaying=            # generated
        ),
    ),
)
# different hover breaks sometimes
# different date stuff break sometimes (wasn't showing me august?)

default_hovertemplate = 'Depth: %{y}, Value: %{x}' # how to get this to be the name

GRAPH_HEIGHT_MIN = 125

# move this?
def is_array(value):
    if isinstance(value, str): return False
    if hasattr(value, "magnitude"):
        return is_array(value.magnitude)
    return hasattr(value, "__len__")

def toTarget(axis):
    return axis[0] + axis[-1]

# So this renderer will have a list of tracks, and their contained axis, but not the data
# Its all position data, not styling data
class Plotly(pzr.Renderer):
    def __init__(self, template=defaults):
        javascript() # double it up
        self.template = copy.deepcopy(template)
        self.xaxis_template = copy.deepcopy(self.template["plotly"]["xaxis_template"])
        del self.template["plotly"]["xaxis_template"]
        self.last_fig = None

    def _hidden(self, themes):
        hidden = themes["hidden"]
        if hidden: themes.pop()
        return hidden

   # TODO, units (find nearest with units too)
    def _process_graph_note(self, note, x0, x1, yref):
        if 'range' not in note:
            raise ValueError("Depth note seems to be missing range")
        # TODO: check numerics
        type = 'line'
        if isinstance(note['range'], tuple): type='rect'
        shape = None
        if type=='line':
            shape = dict(
                type='line',
                xref='paper',
                x0=x0,
                x1=x1,
                yref=yref,
                y0=note['range'],
                y1=note['range'],
                line_width=note['weight'] if 'weight' in note else 1,
                line_dash=note['style'] if 'style' in note else 'dot',
                line_color=note['color'] if 'color' in note else 'black',
            )
        else:
            shape = dict(
                type='rect',
                xref='paper',
                x0=x0,
                x1=x1,
                yref=yref,
                y0=note['range'][0],
                y1=note['range'][1],
                line_width=0,
                fillcolor=note['color'] if 'color' in note else 'lightskyblue',
                layer="below",
                opacity=.5,
            )
        annotation = None
        if not ('show_text' in note and not note['show_text']):
            annotation = dict(
                text=note['text'] if 'text' in note else "",
                xref="paper",
                x=1,
                yref=yref,
                y=note['range'] if type=='line' else note['range'][0],
                yshift=5,
                showarrow=False,
            )
        return shape, annotation

    def get_layout(self, graph, xp=None, **kwargs):
        if not isinstance(graph, pozo.Graph):
            raise TypeError("Layout must be supplied a graph object.")
        if not len(graph):
            raise ValueError("There are no tracks, there is nothing to lay out.")
        show_depth = kwargs.pop("show_depth", True)
        depth_axis_position = kwargs.pop("depth_position", 0)
        override_theme = kwargs.pop("override_theme", None)
        override_theme = kwargs.pop("theme_override", override_theme)
        height = kwargs.get("height", None)
        depth_range = kwargs.get("depth", None)
        if depth_range is not None and ( not isinstance(depth_range, (tuple, list)) or len(depth_range) != 2 ):
            raise TypeError(f"Depth range must be a list or tuple of length two, not {depth_range}")

        layout = copy.deepcopy(self.template["plotly"])
        if height is not None:
            layout["height"] = height
        if layout["height"] < GRAPH_HEIGHT_MIN:
            raise ValueError("125px is the minimum total height")

        tracks_y_axis = "yaxis1"
        xp_y_axis = None
        xp_x_axis = None
        if xp is not None:
            tracks_y_axis = "yaxis2"
            xp_y_axis = "yaxis1"
            xp_x_axis = "xaxis1"
            layout[tracks_y_axis] = copy.deepcopy(layout['yaxis1'])
            del layout['yaxis1']
            # Honestly all communication between renderers should maybe not happen inside any get_layout
            # nor any get_trace

        # TODO constrain positions and domains

        # posmap is an object eventual assigned to the figure containing axis/graph/position data.
        # it's used to pull out reference numbers easily
        posmap = {
            'xp_x': xp_x_axis,
            'xp_y': xp_y_axis,
            'track_y': tracks_y_axis,
            'pixel_height': layout['height'],
            'pixel_width': 4, # start drawing axes at 4 makes it look a little nicer on the left
            'pixel_cursor': 4,
            'with_xp': xp is not None,
            'xp_end': 0,
            'show_depth': show_depth,
            'depth_track_number': depth_axis_position if show_depth else 0, # putting at 0 easier if turned off
            'depth_auto_left': not depth_axis_position, # weird because if `width_xp` this should be false, but logic works
            'depth_auto_right': False, # calc later
            'tracks_axis_numbers': [],
            'tracks_pixel_widths': [],
            'tracks_x_domains': [],
            'tracks_y_domain': None,
            'axis_line_positions': [],
            'total_axes': int(xp is not None),
            'y_bottom': self.template['y_annotation_gutter']/layout['height'],
            }


        # we do a lot of Y calculations in the first run
        max_axis_stack = 0
        track_index = -1
        for track in graph.get_tracks():
            track_index += 1
            if pzt.ThemeStack(track.get_theme())['hidden'] or len(track) == 0: continue
            if track_index and posmap['depth_track_number'] == track_index: posmap['tracks_axis_numbers'].append("depth")
            posmap['tracks_axis_numbers'].append([])
            for axis in track.get_axes():
                if pzt.ThemeStack(track.get_theme())['hidden']: continue
                posmap['total_axes'] += 1 # plotly axis indicies are base 1 so increment first
                posmap['tracks_axis_numbers'][-1].append(posmap['total_axes'])
            max_axis_stack = max(len(posmap['tracks_axis_numbers'][-1]), max_axis_stack)
        if not max_axis_stack: raise ValueError("Didn't find an axes, not going to render")
        if max_axis_stack == 1: posmap['axis_line_positions'] = [1]
        elif max_axis_stack > 1:
            for i in reversed(range(max_axis_stack+1)):
                posmap['axis_line_positions'].append(
                    1 - i * (self.template["axis_label_height"]/posmap['pixel_height'])
                    )
        layout[tracks_y_axis]["domain"] = posmap['tracks_y_domain'] = ( posmap['y_bottom'], posmap['axis_line_positions'][0] )
        if posmap['depth_track_number'] >= len(posmap['tracks_axis_numbers']):
            posmap['depth_auto_right'] = True
        max_text = 0
        for name, note in graph.depth_notes.items():
            max_text = max(max_text, len(note['text'] if 'text' in note else name))
        posmap['x_annotation_pixel_width'] = max_text*10
        posmap['pixel_width'] += max_text*10

        ## in the second pass, we apply units and styles, and information from the first run
        axes_styles = [] # why not declare axes_styles earlier and then apply axes styles from the start, instead of here

        ymin = float('inf')
        ymax = float('-inf')
        y_unit = None

        themes = pzt.ThemeStack(pzt.default_theme, theme = override_theme)
        themes.append(graph.get_theme())
        if self._hidden(themes): return {}

        track_index = -1 # we can't use enumerate() becuase sometimes we skip iterations in the loop
        for track in graph.get_tracks():
            themes.append(track.get_theme())
            if self._hidden(themes) or len(track) == 0: continue
            track_index += 1
            if track_index and posmap['tracks_axis_numbers'][track_index] == 'depth': # don't do it for 0
                track_index +=1
                posmap['tracks_pixel_widths'].append(self.template['depth_axis_width'])


            axis_index = -1 # same story with track_index, need to skip iterations, enumerate() won't work
            for axis in track.get_axes():
                themes.append(axis.get_theme())
                if self._hidden(themes): continue
                axis_index += 1

                if themes["range_unit"] is not None:
                    range_unit = pzu.registry.parse_units(themes["range_unit"])
                else:
                    range_unit = None
                data_unit = None
                for trace in axis:
                    themes.append(trace.get_theme())
                    if self._hidden(themes): continue

                    if data_unit is not None and data_unit != trace.get_unit():
                        raise ValueError(f"Data being displayed on one axis must be exactly the same unit. {data_unit} is not {trace.get_unit()}")
                    else:
                        data_unit = trace.get_unit()
                        if not (data_unit is None or range_unit is None or range_unit.is_compatible_with(data_unit)):
                            raise pint.DimensionalityError(range_unit, data_unit, extra_msg="range_unit set by theme is not compatible with data units")
                    themes.pop()
                    if y_unit is not None and trace.get_depth_unit() is not None:
                        if y_unit != trace.get_depth_unit():
                            raise ValueError(f"All depth axis must have the same unit. You must transform the data. {y_unit} is not {trace.get_depth_unit()}")
                    y_unit = trace.get_depth_unit()
                    ymin = min(pzu.Q(trace.get_depth()[0], y_unit), pzu.Q(ymin, y_unit)).magnitude
                    ymax = max(pzu.Q(trace.get_depth()[-1], y_unit), pzu.Q(ymax, y_unit)).magnitude

                if data_unit is None: data_unit = range_unit
                if range_unit is None: range_unit = data_unit # both None, or both whatever wasn't None
                if data_unit != range_unit:
                    xrange = pzu.Q(themes["range"], range_unit).m_as(data_unit)
                else:
                    xrange = themes["range"]
                # So we've just created xrange which is the data_unit
                # But here we'd want to override ticks and such to the range unit
                color = themes["color"]

                scale_type = themes["scale"]
                axis_style = dict(
                    **self.xaxis_template
                )
                append_unit = ""
                if data_unit is not None:
                    unit = format(data_unit, '~')
                    unit = re_space.sub('', unit)
                    unit = re_power.sub('^', unit)
                    append_unit = " (" + unit + ")"
                axis_style['title'] = dict(text=axis.get_name() + append_unit, font=dict(color=color), standoff=0,)
                axis_style['linecolor'] = color
                axis_style['tickcolor'] = color
                axis_style['tickfont']  = dict(color=color,)

                if axis_style is not None:
                    axis_style['type'] = scale_type
                    if scale_type == "log":
                        xrange = [math.log(xrange[0], 10), math.log(xrange[1], 10)]
                if xrange is not None:
                    axis_style['range'] = xrange

                axis_style['side'] = "top"
                if axis_index: # is a secondary axis
                    axis_style['anchor'] = "free" # TODO we might be able to just change this, check
                    axis_style['overlaying'] = "x" + str(posmap['tracks_axis_numbers'][track_index][0])
                # TODO why do we need position on automatically anchored axis in the case of XP? (from previous comment)
                axis_style['position'] = min(posmap['axis_line_positions'][axis_index], 1)
                axes_styles.append(axis_style)
                themes.pop()
            posmap['pixel_width'] += themes["track_width"] + self.template["track_margin"]
            posmap['tracks_pixel_widths'].append(themes["track_width"])
            themes.pop()
        posmap['pixel_width'] -= self.template["track_margin"] # we don't need margin at the end
        if posmap['depth_track_number'] > 0: posmap['pixel_width'] += self.template['depth_axis_width'] # We'll need space

        # all stuff dependent on position could go here
        if xp is not None:
            posmap["pixel_width"] += layout["height"] # create a square for the crossplot
            if posmap['depth_auto_left']:
                posmap["pixel_width"] += self.template['depth_axis_width']
            posmap['xp_end'] = layout["height"] / (posmap["pixel_width"])
            xp_layout = xp.create_layout(container_width = posmap["pixel_width"], size = layout["height"]) # pre posmap
            layout[xp_y_axis] = xp_layout["yaxis1"]
            layout[xp_x_axis] = xp_layout["xaxis1"]
            layout[xp_x_axis]["domain"] = (0, posmap['xp_end'])
            posmap["xp_domain"] = layout[xp_x_axis]["domain"]
            layout["legend"]["x"] = posmap['xp_end'] * .8
            posmap['pixel_cursor'] += layout["height"]
            if posmap['depth_auto_left']: posmap['pixel_cursor'] += self.template['depth_axis_width']
        track = 0
        for i, axis in enumerate(axes_styles):
            if 'overlaying' in axis:
                axis['domain'] = axes_styles[i-1]['domain']
            else:
                axis['domain'] = (posmap['pixel_cursor']/posmap['pixel_width'],
                                  (posmap['pixel_cursor']+posmap['tracks_pixel_widths'][track])/posmap['pixel_width'])
                posmap['pixel_cursor'] += posmap['tracks_pixel_widths'][track]+self.template["track_margin"]
                posmap['tracks_x_domains'].append(axis['domain'])
                track+= 1
                if ( track < len(posmap['tracks_axis_numbers']) and
                    posmap['tracks_axis_numbers'][track] == 'depth'
                ): # never checks track 0, that's auto
                    track+=1
                    layout[tracks_y_axis]['position'] = (-4 + posmap['pixel_cursor'] + self.template['depth_axis_width'])/posmap['pixel_width']
                    posmap['tracks_x_domains'].append((posmap['pixel_cursor']/posmap['pixel_width'], layout[tracks_y_axis]['position']))
                    posmap['pixel_cursor'] += self.template['depth_axis_width']
            layout["xaxis" + str(i + 1 + int(xp is not None))] = axis
        if not show_depth:
            layout[tracks_y_axis]['showticklabels'] = False
            layout[tracks_y_axis]['ticklen'] = 0
            layout[tracks_y_axis]['ticks'] = ""
        elif posmap['depth_auto_right']:
            layout[tracks_y_axis]['position'] = 1 - posmap['x_annotation_pixel_width']/posmap['pixel_width']
        elif posmap['depth_auto_left'] and xp is not None:
            layout[tracks_y_axis]['position'] = posmap['xp_end'] + self.template['depth_axis_width']/posmap['pixel_width']
        layout[tracks_y_axis]['maxallowed'] = ymax
        layout[tracks_y_axis]['minallowed'] = ymin
        if depth_range is not None:
            layout[tracks_y_axis]['range'] = [depth_range[1], depth_range[0]]
        else:
            layout[tracks_y_axis]['range'] = [ymax, ymin]
        layout['width']=posmap['pixel_width']
        self._last_posmap = posmap


        # TODO: add verts and track depth/ranges, it could be condensed into just passing posmap
        depth_margin = self.template['depth_axis_width']/posmap['pixel_width'] if (posmap['depth_auto_left'] and posmap['with_xp']) else 0
        layout['shapes'] = []
        layout['annotations'] = []
        for name, note in graph.depth_notes.items():
            if 'text' not in note: note['text']=name
            s, a = self._process_graph_note(note, posmap['xp_end']+depth_margin, 1, toTarget(posmap['track_y']))
            if s: layout['shapes'].append(s)
            if a: layout['annotations'].append(a)
        return layout

    def get_traces(self, graph, yaxis='y1', xstart=1, **kwargs):
        override_theme = kwargs.pop("override_theme", None)
        override_theme = kwargs.pop("theme_override", override_theme)
        traces = []
        num_axes = xstart
        themes = pzt.ThemeStack(pzt.default_theme, theme = override_theme)
        themes.append(graph.get_theme())
        if self._hidden(themes): return []
        for track in graph:
            themes.append(track.get_theme())
            if self._hidden(themes): continue
            for axis in track:
                themes.append(axis.get_theme())
                if self._hidden(themes): continue
                all_traces = []
                for trace in axis:
                    themes.append(trace.get_theme())
                    if self._hidden(themes): continue
                    color = themes["color"]
                    with warnings.catch_warnings():
                        warnings.filterwarnings(action='ignore', category=pint.UnitStrippedWarning, append=True)
                        all_traces.append(go.Scattergl(
                            x=trace.get_data(),
                            y=trace.get_depth(),
                            mode='lines', # nope, based on data w/ default
                            line=dict(color=color, width=1), # needs to be better, based on data
                            xaxis='x' + str(num_axes),
                            yaxis=toTarget(yaxis),
                            name = trace.get_name(),
                            hovertemplate = default_hovertemplate,
                            showlegend = False,
                        ))
                    themes.pop()
                num_axes += 1
                traces.extend(all_traces)
                themes.pop()
            themes.pop()
        return traces

    def render(self, graph, static=False, depth=None, xp=None, **kwargs):
        xp_depth = kwargs.pop("xp_depth", None)
        xp_depth_by_index = kwargs.pop("xp_depth_by_index", None)
        depth_lock = False
        if xp_depth or xp_depth_by_index:
            depth_lock = True

        by_index = False
        if xp_depth_by_index:
            if xp_depth: raise ValueError("Specify xp_depth or xp_depth_by_index")
            xp_depth = xp_depth_by_index
            by_index = True
        elif not xp_depth:
            xp_depth = depth

        color_lock = kwargs.pop("color_lock", {})
        # kwargs: theme_override, override_theme (same thing)

        # this generates XP layout too, I don't love, would rather set axes
        # separating XP from here would be good but let's call it techdebt
        layout = self.get_layout(graph, xp=xp, depth=depth, **kwargs)

        # what arguments do we need here now
        traces = self.get_traces(graph, xstart=2 if xp else 1, yaxis='yaxis2' if xp else 'yaxis1', **kwargs)
        if xp is not None:
            traces.extend(
                    xp.create_traces(
                        container_width=layout["width"],
                        depth_range=xp_depth,
                        by_index=by_index,
                        size = layout["height"],
                        static=static,
                        yaxis='y1',
                        color_lock=color_lock
                        )
                    )

        if static:
            self.last_fig = go.Figure(data=traces, layout=layout)
            return self.last_fig
        elif xp is None:
            self.last_fig = xpFigureWidget(data=traces, layout=layout)
            self.last_fig._lead_axis = "yaxis1"
        else:
            self.last_fig = xpFigureWidget(data=traces, layout=layout, renderer=xp, depth_lock=depth_lock)
            self.last_fig._lead_axis = "yaxis2"
            self.last_fig.link_depth_to(self.last_fig)
            xp.add_figure(self.last_fig)
        javascript() # double it up
        return self.last_fig

# Use this for all FigureWidgets, have way to return FigureWidget
class xpFigureWidget(go.FigureWidget):
    # pass the controlling graph here, and it will add a callback
    def link_depth_to(self, fig):
        if not isinstance(fig, go.FigureWidget):
            raise TypeError("Supplied fig argument bust be a go.FigureWidget ot have access to interactivity")
        fig.layout.on_change(self._depth_change_cb, fig._lead_axis+'.range', append=True)

    def __init__(self, data=None, layout=None, frames=None, skip_invalid=False, depth_lock=False, **kwargs):
        self._xp_renderer = kwargs.pop("renderer", None)
        if self._xp_renderer:
            self._depth_lock = True
        super().__init__(data=data, layout=layout, frames=frames, skip_invalid=skip_invalid, **kwargs)

    def _depth_change_cb(self, layout, new_range):
        if not self._depth_lock:
            self.set_depth_range(new_range)
        self._tracks_depth_range = new_range

    def lock_depth_range(self):
        if not self._xp_renderer: raise TypeError("lock_depth_range only applies to graphs with a crossplot")
        self._depth_lock = True

    def unlock_depth_range(self):
        if not self._xp_renderer: raise TypeError("unlock_depth_range only applies to graphs with a crossplot")
        self._depth_lock = False
        self.match_depth_range()

    def match_depth_range(self):
        if not self._xp_renderer: raise TypeError("match_depth_range only applies to graphs with a crossplot")
        self.set_depth_range(depth_range=self._tracks_depth_range)

    def set_depth_range(self, depth_range=None):
        if not self._xp_renderer: raise TypeError("set_depth_range only applies to graphs with a crossplot")
        depth_range = sorted(depth_range) if depth_range else None
        color_range_queue=[]
        with self.batch_update():
            for trace in self['data']:
                if not trace.meta or trace.meta.get("filter", None) != 'depth': continue
                trace.x = self._xp_renderer.x.get_data(slice_by_depth=depth_range)
                trace.y = self._xp_renderer.y.get_data(slice_by_depth=depth_range)
                color_data = trace.meta.get('color_data', None)
                if color_data == 'depth': # if xp were to change x or y, this would be wrong in renderer
                    trace.marker.color = self._xp_renderer.x.get_depth(slice_by_depth=(depth_range))
                elif color_data in self._xp_renderer._colors_by_id:
                    trace.marker.color = self._xp_renderer._colors_by_id[color_data].get_data(slice_by_depth=depth_range)
                if 'color_range' in trace.meta and trace.meta['color_range'] is not None: # no need for (None)?
                    color_range_queue.append((trace.name, trace.meta['color_range']))
        with self.batch_update(): # instead of two batches, color_range could take the new info as arguments
            for item in color_range_queue:
                self.set_color_range(item[0], color_range=item[1])

    # we have to set self._color_range so that can recalculate it when depth changes
    def set_color_range(self, name, color_range=(None), auto=False, lock=False):
        if not self._xp_renderer: raise TypeError("set_color_range only applies to graphs with a crossplot")
        for trace in self['data']:
            if trace.meta and 'color_data' in trace.meta:
                if trace.name != name: continue
                if color_range: color_range = list(sorted(color_range))
                calc_key = 'colorscale_calculated'
                select_key = 'colorscale_selected'

                # if we never calculated or what we have is different than what we calculaed
                if (calc_key not in trace.meta
                    or trace.meta[calc_key] != trace.marker.colorscale
                    ):
                    trace.meta[select_key] = trace.marker.colorscale # they must have selected that scale

                colorscale_selected = trace.meta[select_key]
                if auto or color_range==[None]:
                    trace.meta['color_range'] = (None)
                    if calc_key in trace.meta: del trace.meta[calc_key]
                    if select_key in trace.meta: del trace.meta[select_key]
                    trace.marker.colorscale = colorscale_selected
                    return

                data_min = np.nanmin(trace.marker.color)
                data_max = np.nanmax(trace.marker.color)
                if lock:
                    color_range = [data_min, data_max]
                trace.meta['color_range'] = color_range
                if color_range[0] is None: color_range[0] = data_min
                if color_range[1] is None: color_range[1] = data_max

                normalized_pairs = self._xp_renderer.project_color_scale(data_max, data_min, color_range, colorscale_selected)

                trace.marker.colorscale = tuple(normalized_pairs)
                trace.meta[calc_key] = tuple(normalized_pairs)
                # should ctime be the min max for the whole thing

class CrossPlot():
    marker_symbols = ["circle", "diamond", "square", "cross", "x", "pentagon", "start", "hexagram", "starsquare"]
    default_marker = {
            "opacity": .8,
            "size": 5
            }
    def _get_marker_no_color(self):
        marker = self.default_marker.copy()
        marker["symbol"] = self.marker_symbols[len(self.marker_symbols) % self._marker_symbol_index]
        self._marker_symbol_index += 1
        return marker

    def _get_marker_with_color(self, array, title=None, colorscale="Viridis_r", container_width=None, size=None):
        if not size: size = self.size
        marker = self._get_marker_no_color()
        marker["color"] = array
        marker["showscale"] = True
        marker["colorscale"] = colorscale
        if container_width is not None:
            marker["colorbar"] = dict( # again, problem here, this should be in layout, but its in trace, and that sucks
                    title=title,
                    orientation='h',
                    thickness=20,
                    thicknessmode='pixels',
                    x=(self.size/(2.00)-45)/container_width,
                    xref='paper',
                    y=10/size, # 10% margin, does this really work for all sizes?
                    yref='paper',
                    len=size*.9,
                    lenmode='pixels')
        else:
            marker["colorbar"] = dict(
                    title=title,
                    orientation='h',
                    thickness=20,
                    thicknessmode='pixels')
        marker["opacity"] = self.fade
        return marker

    def _resolve_selector_to_data(self, selector):
        POZO_OBJS = (pozo.Graph, pozo.Track, pozo.Axis)
        if isinstance(selector, POZO_OBJS):
           data = selector.get_data()
           if len(data) == 0:
               raise ValueError(f"{selector} has no pozo.Trace object")
           return data[0] # we process it in the following
        elif isinstance(selector, pozo.Trace):
            if not is_array(selector.get_data()): raise TypeError(f"{selector}'s data seems weird: {selector.get_data()}")
            if not is_array(selector.get_depth()):raise TypeError(f"{selector}'s depth seems weird: {selector.get_depth()}")
            return selector
        raise TypeError(f"{selector} does not appear to be a pozo object, it's a {type(selector)}")

    def reinit(self, x = None, y = None, colors=[None], **kwargs):
        self.size                = kwargs.pop("size", self.size)
        self.depth_range         = kwargs.pop("depth_range", self.depth_range)
        self.xrange              = kwargs.pop("xrange", self.xrange)
        self.yrange              = kwargs.pop("yrange", self.yrange)
        if not is_array(colors): colors = [colors]
        self.colors = colors
        self.y              = y if y is not None else self.y
        self.x              = x if x is not None else self.x

    def __init__(self, x=None, y=None, colors=[None], **kwargs):
        # rendering defaults
        self.size                = kwargs.pop("size", 500)
        self.depth_range         = kwargs.pop("depth_range", [None])
        self.xrange              = kwargs.pop("xrange", None)
        self.yrange              = kwargs.pop("yrange", None)
        self.fade = 1
        if not is_array(colors): colors = [colors]
        self.colors = colors
        self.y = y
        self.x = x
        self._colors_by_id = {}
        self._figures_by_id = weakref.WeakValueDictionary() # Do figures contain their colors? so we can clear up _colors_by_id

    def add_figure(self, fig):
        self._figures_by_id[id(fig)] = fig

    def render(self, **kwargs):
        static = kwargs.pop("static", False)
        layout = self.create_layout(**kwargs) # container_width, size
        traces = self.create_traces(**kwargs) # container_width, depth_range, size, static

        if static:
            self.last_fig = go.Figure(data=traces, layout=layout)
            return self.last_fig
        fig = xpFigureWidget(data=traces, layout=layout, renderer=self)
        self.add_figure(fig)
        self.last_fig = fig

        return fig

    def create_layout(self, container_width=None, size=None, xaxis="xaxis1", yaxis="yaxis1"):
        if not size: size = self.size
        margin = (120) / size if container_width is not None else 0
        return {
            "width"       : size,
            "height"      : size,
            xaxis       : dict(
                            title = self.x.get_name(),
                            range = self.xrange,
                            linecolor = "#888",
                            linewidth = 1,
            ),
            yaxis       : dict(
                            title = self.y.get_name(),
                            range = self.yrange,
                            domain = (margin, 1),
                            linecolor = "#888",
                            linewidth = 1,
            ),
            "showlegend"  : True
        }

    def create_traces(self, depth_range=None, container_width=None, size=None, static=False, xaxis="xaxis1", yaxis="yaxis1", color_lock={}, by_index=False):
        if not size: size = self.size
        if not depth_range: depth_range = self.depth_range
        x_data = self.x.get_data()[slice(*depth_range)] if by_index else self.x.get_data(slice_by_depth=depth_range)
        y_data = self.y.get_data()[slice(*depth_range)] if by_index else self.y.get_data(slice_by_depth=depth_range)

        self._marker_symbol_index = 1

        # Doing some stats
        missing = (np.isnan(x_data) + np.isnan(y_data)).sum() # noqa TODO
        # TODO improve with more statistics, how many values are there, etc
        title = "Number of unplottable values: {missing} ({(100 * missing/len(x_data)):.1f}%)" # noqa TODO
        # title will be added but also updated whenever written so maybe this should be a function

        self._base_trace = dict(
            x = x_data,
            y = y_data,
            mode='markers',
            xaxis=toTarget(xaxis),
            yaxis=toTarget(yaxis),
        )

        # Make Traces
        trace_definitions = []
        plotly_traces     = []
        for color in self.colors:
            trace_definitions.append(self.create_trace(color, container_width=container_width, depth_range=depth_range, size=size, static=static, by_index=by_index))
        if trace_definitions and 'visible' in trace_definitions[0]: del trace_definitions[0]['visible']

        for trace in trace_definitions:
            plotly_traces.append(go.Scattergl(trace))
            if trace['name'] in color_lock:
                # originally we did this post-process modification  because we didn't have plotly.colors.get_colorscale(str)
                # we could now do it in create trace
                color_range = color_lock[trace['name']]
                cs_sel = plotly_traces[-1]['marker']['colorscale'] # this is the selected
                data_min = np.nanmin(plotly_traces[-1]['marker']['color'])
                data_max = np.nanmax(plotly_traces[-1]['marker']['color'])
                cs_calc = self.project_color_scale(data_max, data_min, color_range, cs_sel)
                plotly_traces[-1]['marker']['colorscale'] = cs_calc
                plotly_traces[-1].meta['color_range'] = tuple(color_range)
                plotly_traces[-1].meta['colorscale_calculated'] = tuple(cs_calc)
                plotly_traces[-1].meta['colorscale_selected'] = tuple(cs_sel)
                # this should lock it, otherwise, it goes back to auto
        return plotly_traces


    def create_trace(self, color, container_width=None, depth_range=None, size=None, static=False, by_index=False):
        if not size: size = self.size
        trace = self._base_trace.copy()
        trace['meta']={'filter':'depth'}
        if color is not None:
            if isinstance(color, str) and color.lower() == "depth":
                trace['meta']['color_data'] = 'depth'
                color_array = self.x.get_depth()[slice(*depth_range)] if by_index else self.x.get_depth(slice_by_depth=depth_range)
            else:
                color_data = self._resolve_selector_to_data(color)
                color_array = color_data.get_data()[slice(*depth_range)] if by_index else color_data.get_data(slice_by_depth=depth_range)
                if not static: # no need to store if image will never update
                    trace['meta']['color_data'] = id(color)
                    self._colors_by_id[id(color)] = color

            name = color_data.get_name() if color != "depth" else "depth"
            trace['name'] = name
            trace['marker'] = self._get_marker_with_color(color_array, name, container_width=container_width, size=size)
            trace['hovertemplate'] = '%{x}, %{y}, %{marker.color}'
            trace['visible'] = 'legendonly'
        else:
            trace['marker'] = self._get_marker_no_color()
            trace['name'] = "x"
        return trace

    def project_color_scale(self, data_max, data_min, color_range, colorscale_selected):
        distance = data_max - data_min
        normalized_color_range = np.round((color_range - data_min) / distance, 2)
        norm_distance = normalized_color_range[1] - normalized_color_range[0]

        normalized_pairs = [] # lets create normalized pairs
        for pair in colorscale_selected:
            normalized_pairs.append(((pair[0]*norm_distance)+normalized_color_range[0], pair[1]))
        if normalized_pairs[0][0] > 1: normalized_pairs = [(0, normalized_pairs[0][1]), (1,normalized_pairs[0][1]) ]
        elif normalized_pairs[-1][0] < 0: normalized_pairs = [(0, normalized_pairs[-1][1]), (1,normalized_pairs[-1][1]) ]
        else:
            # can bottom be extended
            if normalized_pairs[0][0] > 0: normalized_pairs.insert(0, (0, normalized_pairs[0][1]))
            elif normalized_pairs[0][0] < 0:
                lb = normalized_pairs[0]
                ub = None
                old_pairs = normalized_pairs.copy()
                for pair in old_pairs:
                    if pair[0] > 0:
                        if ub is None:
                            ub = pair
                            #print(lb)
                            new_color = plotly.colors.label_rgb(
                                plotly.colors.find_intermediate_color(
                                    plotly.colors.hex_to_rgb(lb[1]) if lb[1][0] != 'r' else plotly.colors.unlabel_rgb(lb[1]),
                                    plotly.colors.hex_to_rgb(ub[1]) if ub[1][0] != 'r' else plotly.colors.unlabel_rgb(ub[1]),
                                    (-lb[0])/(ub[0]-lb[0]) # distance from 0
                                    ))
                            #print(ub)
                            normalized_pairs = [(0, new_color)]
                        normalized_pairs.append(pair)
                    else:
                        lb = pair
            if normalized_pairs[-1][0] < 1: normalized_pairs.append((1, normalized_pairs[-1][1]))
            elif normalized_pairs[-1][0] > 1:
                ub = normalized_pairs[-1]
                lb = None
                old_pairs = list(reversed(normalized_pairs.copy()))
                for pair in old_pairs:
                    if pair[0] < 1: # we're here
                        if lb is None:
                            lb = pair
                            #print(lb)
                            #print(ub)
                            new_color = plotly.colors.label_rgb(
                                plotly.colors.find_intermediate_color(
                                    plotly.colors.hex_to_rgb(lb[1]) if lb[1][0] != 'r' else plotly.colors.unlabel_rgb(lb[1]),
                                    plotly.colors.hex_to_rgb(ub[1]) if ub[1][0] != 'r' else plotly.colors.unlabel_rgb(ub[1]),
                                    (1-lb[0])/(ub[0]-lb[0]) # distance from 0
                                    ))
                            normalized_pairs = [(1, new_color)]
                        normalized_pairs.append(pair)
                    else:
                        ub = pair
                normalized_pairs = list(reversed(normalized_pairs))
        return normalized_pairs

    @property
    def colors(self):
        return self.__colors

    @colors.setter
    def colors(self, colors):
        self.__colors = []
        for color in colors:
            if color is None:
                self.__colors.append(None)
            elif isinstance(color, str) and color.lower() == "depth":
                self.__colors.append("depth")
            else:
                self.__colors.append(self._resolve_selector_to_data(color))

    @property
    def x(self):
        return self.__x
    @x.setter
    def x(self, x):
        self.__x = self._resolve_selector_to_data(x) if x is not None else None

    @property
    def y(self):
        return self.__y
    @y.setter
    def y(self, y):
        self.__y = self._resolve_selector_to_data(y) if y is not None else None


def init_write_image(counter):
    global write_image_counter
    write_image_counter = counter

def write_image(gp):
    gp['graph'].write_image(gp['path'], engine="kaleido")
    with write_image_counter.get_lock():
        write_image_counter.value += 1

def make_xp_depth_video(folder_name, graph, start, window, end, xp=True, output="all.mp4", fps=30):
    try:
        write_image_counter
        raise Exception("Please only run make_frames once at a time, or restart the kernel")
    except NameError:
        pass
    if output:
        try:
            import imageio
        except ImportError:
            raise ImportError("Please install imageio. It must be installed like: pip install imageio[ffmpeg] or pip install imageio[pyav]")
    os.makedirs(folder_name, exist_ok=True)
    if xp is True: xp=graph.xp
    elif not isinstance(xp, CrossPlot): raise ValueError("Crossplot renderer not valid. xp=True (default) or another renderer)")
    trace = xp.x
    depth = trace.get_depth()
    start_index = trace.find_nearest(start)[0]
    window_index = trace.find_nearest(start+window)[0] - start_index
    end_index = trace.find_nearest(end)[0]
    frame_count = range(start_index, end_index-window_index, 1)
    gp = []
    render_counter = IntProgress(min=0, max=len(frame_count)-1, description="Rendering:")
    writer_counter = IntProgress(min=0, max=len(frame_count)-1, description="Writing:")
    display(render_counter)
    display(writer_counter)
    tail_size = math.floor(window_index*.2)
    tail = np.linspace(0,1,tail_size)
    fade = tail.tolist() + [1]*(window_index-tail_size)
    for i, cursor in enumerate(frame_count):
        render_counter.value += 1
        graph.depth_notes['Depth Highlight-xxx'] = dict(range = (depth[cursor], depth[cursor+window_index]), show_text=False)
        graph.xp.fade = fade
        gp.append(
                dict(
                    graph=graph.render(
                        height=800,
                        xp=True,
                        depth_position=1,
                        color_lock={'depth': [start, end]},
                        depth=[start, end],
                        xp_depth_by_index=[cursor, cursor+window_index],
                        static=True,
                        ),
                    path=folder_name+"/"+str(i)+".png",
                    )
                )
    del graph.depth_notes['Depth Highlight-xxx']
    write_image_counter = multiprocessing.Value('I',0)
    with multiprocessing.Pool(initializer=init_write_image, initargs=(write_image_counter,)) as pool:
        p = pool.map_async(write_image, gp)
        while True:
            p.wait(.5)
            try:
                p.successful()
                break
            except ValueError:
                writer_counter.value = write_image_counter.value
    writer_counter.value = len(frame_count)-1
    del write_image_counter
    if not output: return
    ims = [imageio.imread(p['path']) for p in gp]
    imageio.mimwrite(folder_name+"/" + output, ims, fps=fps)
