# -*- encoding: utf-8 -*-

# @Author : soarkey
# @Date   : 2022/3/7
# @Desc   : block-based cooperate attention module

import torch
import torch.nn as nn
import torch.nn.functional as F

from lib.core.config import config
from lib.models.dcn import DeformConvPack

# channel attention
class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            nn.Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(inplace=True),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
        )
        self.pool_types = pool_types

    def forward(self, x):
        """
        Args:
            x: [B,256,7,7]
        Returns:
            c_weight: [B,256,1,1]
        """
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type == 'avg':
                avg_pool = F.avg_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp(avg_pool)
            elif pool_type == 'max':
                max_pool = F.max_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp(max_pool)
            elif pool_type == 'lp':
                lp_pool = F.lp_pool2d(x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp(lp_pool)

            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw

        scale = torch.sigmoid(channel_att_sum).unsqueeze(2).unsqueeze(3)
        return scale

# block-based cooperate attention
class BCA(nn.Module):
    def __init__(self, in_channel=256):
        super(BCA, self).__init__()

        self.in_channel = in_channel
        self.hidden_channel = in_channel // 2

        # feature extraction
        if config.TRAIN.ATTN.BLOCK_FEATURE:
            # template branch
            self.g1 = nn.Sequential(
                DeformConvPack(in_channel, self.hidden_channel, kernel_size=3, stride=1, padding=0),
                nn.BatchNorm2d(self.hidden_channel),
                nn.ReLU(inplace=True),
                DeformConvPack(self.hidden_channel, self.hidden_channel, kernel_size=3, stride=1, padding=0),
                nn.BatchNorm2d(self.hidden_channel),
                nn.ReLU(inplace=True),
            )

            # search branch
            self.g2 = nn.Sequential(
                DeformConvPack(in_channel, self.hidden_channel, kernel_size=3, stride=2, padding=0),
                nn.BatchNorm2d(self.hidden_channel),
                nn.ReLU(inplace=True),
            )

            self.theta = DeformConvPack(in_channel, self.hidden_channel, kernel_size=3, stride=2, padding=0)
            self.phi = DeformConvPack(in_channel, self.hidden_channel, kernel_size=3, stride=2, padding=0)

            self.Q = nn.Sequential(
                DeformConvPack(self.hidden_channel, 2 * in_channel, kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(2 * in_channel),
                nn.ReLU(inplace=True),
                nn.PixelShuffle(2),
                nn.UpsamplingBilinear2d(size=(7, 7)),
                DeformConvPack(self.hidden_channel, in_channel, kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(in_channel),
                nn.ReLU(inplace=True),
            )

            self.K1 = nn.Sequential(
                DeformConvPack(self.hidden_channel, 2 * in_channel, kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(2 * in_channel),
                nn.ReLU(inplace=True),
                nn.PixelShuffle(upscale_factor=2),
            )

            self.K2 = nn.Sequential(
                DeformConvPack(self.hidden_channel, in_channel, kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(in_channel),
                nn.ReLU(inplace=True),
            )

            self.alpha = nn.Parameter(1 * torch.ones(1).cuda())
            self.beta = nn.Parameter(1 * torch.ones(1).cuda())

        if config.TRAIN.ATTN.CHANNEL_ATTN:
            # channel attention
            self.ChannelGate = ChannelGate(in_channel)

    def forward(self, z, x):
        """
        args:
            z: [B,256,7,7]
            x: [B,256,31,31]
        returns:
            z: [B,256,7,7]
            x: [B,256,31,31]
        """

        b, _, z_height, z_width = z.shape
        _, _, x_height, x_width = x.shape

        if config.TRAIN.ATTN.BLOCK_FEATURE:
            zh, zw, xh, xw = [(x - 2) // 2 + 1 for x in [z_height, z_width, x_height, x_width]]

            zg = self.g1(z).view(b, self.hidden_channel, -1)
            zg = zg.permute(0, 2, 1).contiguous()

            xg = self.g2(x).view(b, self.hidden_channel, -1)
            xg = xg.permute(0, 2, 1).contiguous()

            z_theta = self.theta(z).view(b, self.hidden_channel, -1)
            z_theta = z_theta.permute(0, 2, 1)

            x_phi = self.phi(x).view(b, self.hidden_channel, -1)

            f = torch.matmul(z_theta, x_phi)  # [32,25,64]*[32,64,625]=[32,25,625]
            n = f.size(-1)
            fz_div_c = f / n

            f = f.permute(0, 2, 1).contiguous()
            n = f.size(-1)
            fx_div_c = f / n

            zn = torch.matmul(fz_div_c, xg)
            zn = zn.permute(0, 2, 1).contiguous().view(b, self.hidden_channel, zh, zw)
            zn = self.Q(zn)

            zq = self.alpha * zn + z

            xn = torch.matmul(fx_div_c, zg)
            xn = xn.permute(0, 2, 1).contiguous().view(b, self.hidden_channel, xh, xw)
            xn = self.K2(F.upsample_bilinear(self.K1(xn), size=(x_width, x_height)))  # [32,256,31,31]

            xk = self.beta * xn + x
        else:
            zq = z
            xk = x

        # response in channel weight
        if config.TRAIN.ATTN.CHANNEL_ATTN:
            # channel attention
            c_weight = self.ChannelGate(zq)
            zq = zq * c_weight
            xk = xk * c_weight

        return zq, xk

if __name__ == '__main__':
    kernel = torch.randn([32, 256, 7, 7]).cuda()
    search = torch.randn([32, 256, 31, 31]).cuda()

    # params & latency
    from thop.profile import profile
    from thop.utils import clever_format
    import time

    def calc_macs_and_params(name, model, kernel, search):
        model = model.cuda()
        macs, params = profile(model, inputs=(kernel, search), verbose=False)
        macs, params = clever_format([macs, params], "%3.3f")

        # test latency
        T_w = 20
        T_t = 100
        with torch.no_grad():
            for i in range(T_w):
                _, _ = model(kernel, search)
            start = time.time()
            for i in range(T_t):
                _, _ = model(kernel, search)
            end = time.time()
            latency = ((end - start) / T_t) * 1000

        print(f"{name.ljust(25)} | macs {macs} | params {params} | latency {latency}")

    calc_macs_and_params('bca', BCA(), kernel, search)


"""Edit Colors dialog."""

from typing import Any, Callable

from rich.segment import Segment
from rich.style import Style
from textual import events
from textual.color import Color as Color
from textual.containers import Container, Horizontal, Vertical
from textual.css.query import NoMatches
from textual.geometry import Offset
from textual.message import Message
from textual.reactive import reactive, var
from textual.strip import Strip
from textual.widget import Widget
from textual.widgets import Button, Input, Label

from textual_paint.args import args
from textual_paint.localization.i18n import get as _
from textual_paint.windows import DialogWindow

# https://github.com/kouzhudong/win2k/blob/ce6323f76d5cd7d136b74427dad8f94ee4c389d2/trunk/private/shell/win16/comdlg/color.c#L38-L43
# These are a fallback in case colors are not received from some driver.
# const default_basic_colors = [
# 	"#8080FF", "#80FFFF", "#80FF80", "#80FF00", "#FFFF80", "#FF8000", "#C080FF", "#FF80FF",
# 	"#0000FF", "#00FFFF", "#00FF80", "#40FF00", "#FFFF00", "#C08000", "#C08080", "#FF00FF",
# 	"#404080", "#4080FF", "#00FF00", "#808000", "#804000", "#FF8080", "#400080", "#8000FF",
# 	"#000080", "#0080FF", "#008000", "#408000", "#FF0000", "#A00000", "#800080", "#FF0080",
# 	"#000040", "#004080", "#004000", "#404000", "#800000", "#400000", "#400040", "#800040",
# 	"#000000", "#008080", "#408080", "#808080", "#808040", "#C0C0C0", "#400040", "#FFFFFF",
# ];
# Grabbed with Color Cop from the screen with Windows 98 SE running in VMWare
basic_colors = [
	"#FF8080", "#FFFF80", "#80FF80", "#00FF80", "#80FFFF", "#0080FF", "#FF80C0", "#FF80FF",
	"#FF0000", "#FFFF00", "#80FF00", "#00FF40", "#00FFFF", "#0080C0", "#8080C0", "#FF00FF",
	"#804040", "#FF8040", "#00FF00", "#008080", "#004080", "#8080FF", "#800040", "#FF0080",
	"#800000", "#FF8000", "#008000", "#008040", "#0000FF", "#0000A0", "#800080", "#8000FF",
	"#400000", "#804000", "#004000", "#004040", "#000080", "#000040", "#400040", "#400080",
	"#000000", "#808000", "#808040", "#808080", "#408080", "#C0C0C0", "#400040", "#FFFFFF",
]
custom_colors = [
	"#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF",
	"#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF", "#FFFFFF",
]

num_colors_per_row = 8

class ColorGrid(Container):
    """A grid of colors."""

    class Changed(Message):
        """A message that is sent when the selected color changes."""

        def __init__(self, color: str, color_grid: "ColorGrid", index: int) -> None:
            """Initialize the message."""
            super().__init__()
            self.color = color
            self.color_grid = color_grid
            self.index = index

    color_list: var[list[str]] = var(list[str], init=False)
    """The list of colors to display. NOT TO BE CONFUSED WITH `colors` defined by `Widget`."""

    def __init__(self, color_list: list[str], selected_color: str, **kwargs: Any) -> None:
        """Initialize the ColorGrid."""
        super().__init__(**kwargs)
        self.selected_color: str = selected_color
        self._color_by_button: dict[Button, str] = {}
        self.color_list = color_list  # This immediately calls `watch_color_list`.
        self.can_focus = True

    def on_mount(self) -> None:
        """Called when the window is mounted."""
        found_match = False
        for button, color in self._color_by_button.items():
            matches = Color.parse(color) == Color.parse(self.selected_color)
            if matches and not found_match:
                button.add_class("focused")
                found_match = True
        self._select_focused_color()

    def watch_color_list(self, color_list: list[str]) -> None:
        """Called when the color list changes."""
        self._color_by_button = {}
        for button in self.query(Button):
            button.remove()
        for color in self.color_list:
            button = Button("", classes="color_button color_well")
            button.styles.background = color
            button.can_focus = False  # using fake focus for now
            self._color_by_button[button] = color
            self.mount(button)

    def on_key(self, event: events.Key) -> None:
        """Called when a key is pressed."""
        if event.key == "up":
            self._navigate_relative(-num_colors_per_row)
        elif event.key == "down":
            self._navigate_relative(+num_colors_per_row)
        elif event.key == "left":
            self._navigate_relative(-1)
        elif event.key == "right":
            self._navigate_relative(+1)
        elif event.key == "home":
            self._navigate_absolute(0)
        elif event.key == "end":
            self._navigate_absolute(len(self.color_list) - 1)
        elif event.key in ("space", "enter"):
            self._select_focused_color()

    def _select_focused_color(self) -> None:
        try:
            focused = self.query_one(".focused", Button)
        except NoMatches:
            return
        for selected in self.query(".selected"):
            selected.remove_class("selected")
        focused.add_class("selected")
        self.selected_color = self._color_by_button[focused]
        index = list(self._color_by_button.keys()).index(focused)
        self.post_message(self.Changed(self.selected_color, self, index))

    def _navigate_relative(self, delta: int) -> None:
        """Navigate to a color relative to the currently focused color."""
        try:
            focused = self.query_one(".focused", Button)
        except NoMatches:
            return
        # index = self._colors.index(self._color_by_button[focused]) # doesn't work because there can be duplicates
        index = list(self._color_by_button.keys()).index(focused)
        # print(delta, (index % num_colors_per_row), num_colors_per_row)
        if delta == -1 and (index % num_colors_per_row) == 0:
            return
        if delta == +1 and (index % num_colors_per_row) == num_colors_per_row - 1:
            return
        self._navigate_absolute(index + delta)

    def _navigate_absolute(self, index: int) -> None:
        """Navigate to the color at the given index."""
        if index < 0 or index >= len(self.color_list):
            return
        target_button = list(self._color_by_button.keys())[index]
        target_button.add_class("focused")
        for button in self._color_by_button:
            button.remove_class("focused")
        target_button.add_class("focused")

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Called when a button is clicked or activated with the keyboard."""
        self.selected_color = self._color_by_button[event.button]
        for button in self._color_by_button:
            button.remove_class("focused")
        event.button.add_class("focused")
        self._select_focused_color()
        self.focus()

    # I want MouseDown rather than Pressed in order to implement double-clicking.
    # However, event.control is None for mouse events, so this doesn't work:
    # def on_mouse_down(self, event: events.MouseDown) -> None:
    #     """Called when the mouse is pressed down."""
    #     if event.button == 1:
    #         self.selected_color = self._color_by_button[event.control]
    #         self.refresh()
    # @on(events.MouseDown, ".color_button")
    # def handle_color_button(self, event: events.MouseDown) -> None:
    #     """Called when a color button is clicked."""
    #     self.selected_color = self._color_by_button[event.control]

class LuminosityRamp(Widget):
    """A vertical slider to select a luminosity, previewing the color at each luminosity with a gradient."""

    class Changed(Message):
        """A message that is sent when the luminosity changes."""

        def __init__(self, luminosity: float) -> None:
            """Initialize the Changed message."""
            super().__init__()
            self.luminosity = luminosity

    hue = reactive(0.0)
    saturation = reactive(0.0)
    luminosity = reactive(0.0)

    def __init__(self, hue: float, saturation: float, luminosity: float, **kwargs: Any) -> None:
        """Initialize the LuminosityRamp."""
        super().__init__(**kwargs)
        self.luminosity = luminosity
        self.hue = hue
        self.saturation = saturation
        self._mouse_down = False

    def render_line(self, y: int) -> Strip:
        """Render a line of the widget."""
        marker = "<" if args.ascii_only else "◀" # ◀ (bigger/clearer) or 🢐 (closer shape but smaller)
        # lum = y / self.size.height # bottom isn't quite white
        lum = 1 - y / (self.size.height - 1)
        color = Color.from_hsl(self.hue, self.saturation, lum)
        style = Style(bgcolor=color.rich_color)
        segments = [Segment(" " * (self.size.width - 1), style, None)]
        if y == round((self.size.height - 1) * (1 - self.luminosity)):
            segments.append(Segment(marker, Style(color="black"), None))
        return Strip(segments)

    def on_mouse_down(self, event: events.MouseDown) -> None:
        """Called when the mouse is pressed down."""
        self._update_color(event.y)
        self._mouse_down = True
        self.capture_mouse()

    def on_mouse_up(self, event: events.MouseUp) -> None:
        """Called when the mouse is released."""
        self.release_mouse()
        self._mouse_down = False

    def on_mouse_move(self, event: events.MouseMove) -> None:
        """Called when the mouse is moved."""
        if self._mouse_down:
            self._update_color(event.y)

    def _update_color(self, y: int) -> None:
        """Update the color based on the given y coordinate."""
        self.luminosity = max(0, min(1, 1 - y / (self.size.height - 1)))
        self.post_message(self.Changed(luminosity=self.luminosity))
        # self.refresh()

class ColorField(Widget):
    """A field of hue and saturation, where you can pick a color by clicking."""

    class Changed(Message):
        """A message that is sent when the color changes."""

        def __init__(self, hue: float, saturation: float) -> None:
            """Initialize the Changed message."""
            super().__init__()
            self.hue = hue
            self.saturation = saturation

    hue = reactive(0.0)
    saturation = reactive(0.0)

    def __init__(self, hue: float, saturation: float, **kwargs: Any) -> None:
        """Initialize the ColorField."""
        super().__init__(**kwargs)
        self.hue = hue
        self.saturation = saturation
        self._mouse_down = False

    def render_line(self, y: int) -> Strip:
        """Render a line of the widget."""
        segments: list[Segment] = []
        crosshair = "+" if args.ascii_only else "✜" # Options: ┼+✜✛⊹✚╋╬⁘⁛⌖⯐ or
        #  ╻
        # ╺ ╸
        #  ╹
        for x in range(self.size.width):
            crosshair_here = (
                x == round(self.hue * (self.size.width - 1)) and
                y == round((self.size.height - 1) * (1 - self.saturation))
            )
            color = Color.from_hsl(x / (self.size.width - 1), 1 - y / (self.size.height - 1), 0.5)
            char = crosshair if crosshair_here else " "
            segments.append(Segment(char, Style(color="black", bgcolor=color.rich_color), None))
        return Strip(segments)

    def on_mouse_down(self, event: events.MouseDown) -> None:
        """Called when the mouse is pressed down."""
        self._update_color(event.offset)
        self._mouse_down = True
        self.capture_mouse()

    def on_mouse_up(self, event: events.MouseUp) -> None:
        """Called when the mouse is released."""
        self.release_mouse()
        self._mouse_down = False

    def on_mouse_move(self, event: events.MouseMove) -> None:
        """Called when the mouse is moved."""
        if self._mouse_down:
            self._update_color(event.offset)

    def _update_color(self, offset: Offset) -> None:
        """Update the color based on the given offset."""
        x, y = offset
        self.hue = max(0, min(1, x / (self.size.width - 1)))
        self.saturation = max(0, min(1, 1 - y / (self.size.height - 1)))
        self.post_message(self.Changed(hue=self.hue, saturation=self.saturation))
        # self.refresh()

class ColorPreview(Widget):
    """A preview of the selected color. This doesn't really need to be a special widget..."""

    color: reactive[str] = reactive("black")

    def __init__(self, color: str, **kwargs: Any) -> None:
        """Initialize the ColorPreview."""
        super().__init__(**kwargs)
        self.color = color

    def render_line(self, y: int) -> Strip:
        """Render a line of the widget."""
        return Strip([Segment(" " * self.size.width, Style(bgcolor=self.color), None)])

class IntegerInput(Input):
    """An input that only accepts integers."""

    def __init__(self, min: int, max: int, **kwargs: Any) -> None:
        """Initialize the IntegerInput."""
        super().__init__(**kwargs)
        self.min = min
        self.max = max
        self.last_valid_int = 0

    def _track_valid_int(self, value: str) -> int:
        try:
            value_as_int = int(value)
        except ValueError:
            return self.last_valid_int
        value_as_int = max(self.min, min(self.max, value_as_int))
        self.last_valid_int = value_as_int
        return value_as_int

    def validate_value(self, value: str) -> str:
        """Validate the given value."""
        if not value:
            # Allow empty string
            return value
        return str(self._track_valid_int(value))

    def on_blur(self, event: events.Blur) -> None:
        """Called when the input loses focus. Resets the input if empty."""
        if not self.value:
            self.value = str(self.last_valid_int)

class EditColorsDialogWindow(DialogWindow):
    """A dialog window that lets the user select a color."""

    def __init__(self, *children: Widget, title: str = _("Edit Colors"), selected_color: str|None, handle_selected_color: Callable[[str], None], **kwargs: Any) -> None:
        """Initialize the Edit Colors dialog."""
        super().__init__(handle_button=self.handle_button, *children, title=title, **kwargs)
        self.hue_degrees = 0.0
        self.sat_percent = 0.0
        self.lum_percent = 0.0
        # self._initial_color = selected_color
        if selected_color:
            self._current_color = selected_color
        self._color_by_button: dict[Button, str] = {}
        self._inputs_by_letter: dict[str, IntegerInput] = {}
        self._custom_colors_index = 0
        self.handle_selected_color = handle_selected_color

    def handle_button(self, button: Button) -> None:
        """Called when a button is clicked or activated with the keyboard."""
        if button.has_class("cancel"):
            self.request_close()
        elif button.has_class("ok"):
            self.handle_selected_color(self._current_color.hex)
        elif button.has_class("add_to_custom_colors"):
            global custom_colors
            custom_colors = custom_colors[:]  # copy so that watch_color_list gets called
            # (it uses reference equality; an alternative would be to set always_update=True)
            custom_colors[self._custom_colors_index] = self._current_color.hex
            self.custom_colors_grid.color_list = custom_colors
            self._custom_colors_index = (self._custom_colors_index + 1) % len(custom_colors)

    def on_mount(self) -> None:
        """Called when the window is mounted."""
        self.basic_colors_grid = ColorGrid(basic_colors, self._current_color.hex, classes="autofocus")
        self.custom_colors_grid = ColorGrid(custom_colors, self._current_color.hex)
        verticals_for_inputs: list[Vertical] = []
        for color_model in ["hsl", "rgb"]:
            input_containers: list[Container] = []
            for component_letter in color_model:
                text_with_hotkey: str = {
                    "h": _("Hu&e:"),
                    "s": _("&Sat:"),
                    "l": _("&Lum:"),
                    "r": _("&Red:"),
                    "g": _("&Green:"),
                    "b": _("Bl&ue:"),
                }[component_letter]
                text_without_hotkey = text_with_hotkey.replace("&", "")
                max_value: int = {
                    "h": 360,
                    "s": 100,
                    "l": 100,
                    "r": 255,
                    "g": 255,
                    "b": 255,
                }[component_letter]
                min_value: int = 0
                input = IntegerInput(min_value, max_value, name=component_letter)
                label = Label(text_without_hotkey)
                container = Container(label, input, classes="input_container")
                input_containers.append(container)
                self._inputs_by_letter[component_letter] = input
            verticals_for_inputs.append(Vertical(*input_containers))

        self.content.mount(
            Horizontal(
                Vertical(
                    Label(_("Basic Colors")),
                    self.basic_colors_grid,
                    Label(_("Custom Colors")),
                    self.custom_colors_grid,
                ),
                Vertical(
                    Horizontal(
                        ColorField(self.hue_degrees / 360, self.sat_percent / 100),
                        LuminosityRamp(self.hue_degrees / 360, self.sat_percent / 100, self.lum_percent / 100),
                    ),
                    Horizontal(
                        Vertical(
                            ColorPreview("black"),
                            Label(_("Color")),
                            classes="color_preview_area",
                        ),
                        *verticals_for_inputs,
                    ),
                    Button(_("Add to Custom Colors"), classes="add_to_custom_colors"),
                ),
            ),
            Container(
                Button(_("OK"), classes="ok submit", variant="primary"),
                Button(_("Cancel"), classes="cancel"),
                classes="buttons",
            )
        )

        self._update_inputs("hslrgb")

    def on_color_grid_changed(self, event: ColorGrid.Changed) -> None:
        """Called when the user selects a color from the grid."""
        self._current_color = event.color
        self._update_inputs("hslrgb")
        self._update_color_preview()
        if event.color_grid is self.custom_colors_grid:
            self._custom_colors_index = event.index
        # This is a little awkward, removing .selected for other grids here, but
        # for the clicked grid in the ColorGrid widget itself.
        for color_grid in self.query(ColorGrid):
            if event.color_grid is not color_grid:
                for button in color_grid.query(Button):
                    button.remove_class("selected")

    def on_luminosity_ramp_changed(self, event: LuminosityRamp.Changed) -> None:
        """Called when dragging the luminosity slider."""
        self.lum_percent = event.luminosity * 100
        self._update_inputs("lrgb")
        self._update_color_preview()

    def on_color_field_changed(self, event: ColorField.Changed) -> None:
        """Called when dragging the color field."""
        self.hue_degrees = event.hue * 360
        self.sat_percent = event.saturation * 100
        self._update_inputs("hsrgb")
        self._update_color_preview()

    def on_input_changed(self, event: Input.Changed) -> None:
        """Called when an input changes."""
        component_letter = event.input.name
        if component_letter is None:
            return
        assert isinstance(event.input, IntegerInput)
        value = event.input.last_valid_int
        if component_letter in "hsl":
            if component_letter == "h":
                self.hue_degrees = value
            elif component_letter == "s":
                self.sat_percent = value
            else:
                self.lum_percent = value
            self._update_inputs("rgb")
        else:
            rgb = list(self._current_color.rgb)
            rgb["rgb".index(component_letter)] = value
            self._current_color = Color(*rgb)
            self._update_inputs("hsl")
        self._update_color_preview()

    @property
    def _current_color(self) -> Color:
        """Get the current color."""
        return Color.from_hsl(self.hue_degrees / 360, self.sat_percent / 100, self.lum_percent / 100)

    @_current_color.setter
    def _current_color(self, color: Color | str) -> None:
        """Set the color values from the given textual.color.Color object or string."""
        if isinstance(color, str):
            color = Color.parse(color)
        hue, sat, lum = color.hsl
        self.hue_degrees = hue * 360
        self.sat_percent = sat * 100
        self.lum_percent = lum * 100

    def _update_inputs(self, component_letters: str) -> None:
        """Update the inputs for the given component letters."""
        with self.prevent(Input.Changed):
            for component_letter in component_letters:
                input = self._inputs_by_letter[component_letter]
                input.value = str(int({
                    "h": self.hue_degrees,
                    "s": self.sat_percent,
                    "l": self.lum_percent,
                    "r": self._current_color.rgb[0],
                    "g": self._current_color.rgb[1],
                    "b": self._current_color.rgb[2],
                }[component_letter]))

    def _update_color_preview(self) -> None:
        """Update the color preview."""
        self.query_one(ColorPreview).color = self._current_color.hex
        luminosity_ramp = self.query_one(LuminosityRamp)
        color_field = self.query_one(ColorField)
        luminosity_ramp.luminosity = self.lum_percent / 100
        luminosity_ramp.hue = self.hue_degrees / 360
        luminosity_ramp.saturation = self.sat_percent / 100
        color_field.hue = self.hue_degrees / 360
        color_field.saturation = self.sat_percent / 100


# -*- coding: utf-8 -*-
"""
Created on Sun Oct 27 14:00:50 2019

@author: 朱小五

微信公众号: 凹凸玩数据

哈哈还有一个：凹凸数读
"""

import pandas as pd
import requests
import time
headers = {'User-Agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/77.0.3865.120 Safari/537.36'}

#生成出生当年所有日期
def dateRange(a,b):
    fmt = '%Y-%m-%d'
    bgn = int(time.mktime(time.strptime(a,fmt)))
    end = int(time.mktime(time.strptime(b,fmt)))
    list_date = [time.strftime(fmt,time.localtime(i)) for i in range(bgn,end+1,3600*24)]
    return list_date

def get_json(url):
    try:
        response = requests.get(url,headers=headers)
        if response.status_code == 200:
            json_text=response.json()
            return json_text
    except Exception:
        print('此页有问题！')
        return None

def get_comments(url):
    doc = get_json(url)
    dic = {}
    dic['pigprice'] = doc['pigprice']
    dic['pig_in'] = doc['pig_in']
    dic['pig_local'] = doc['pig_local']
    dic['maizeprice'] = doc['maizeprice']
    dic['bean'] = doc['bean']
    a = '-'.join(doc['time'][3])
    b = time.strftime('%Y-%m-%d',time.localtime(time.time()))
    print(dateRange(a,b))
    dic['time'] = dateRange(a,b)
    return pd.DataFrame(dic)

data =get_comments('http://zhujia.zhuwang.cc/index/api/chartData?areaId=-1&aa=1571997555296')

#作图
from pylab import mpl
import  matplotlib.pyplot as plt
mpl.rcParams['font.sans-serif'] = ['SimHei'] # 指定默认字体
mpl.rcParams['axes.unicode_minus']

plt.figure(figsize=(8,10), dpi=80)
plt.figure(1)
ax1 = plt.subplot(311)
plt.plot(data['time'],data['pigprice'], color="r",linestyle = "-")
plt.xticks([])
plt.annotate(data['pigprice'][365], xy=(data['time'][365], 40), xytext=(data['time'][270], 35), arrowprops=dict(facecolor='black', shrink=0.1, width=0.5))
plt.xlabel("生猪(外三元) 元/公斤")

ax2 = plt.subplot(312)
plt.plot(data['time'],data['maizeprice'],color="y",linestyle = "-")
plt.xticks([])
plt.xlabel("玉米(15%水分) 元/吨")

ax3 = plt.subplot(313)
plt.plot(data['time'],data['bean'],color="g",linestyle = "-")
plt.xlabel("豆粕(43%蛋白) 元/吨")
plt.xticks(data['time'][2::121], rotation=0)


import threading
from django.shortcuts import render, redirect
from django.contrib.contenttypes.models import ContentType
from .models import Comment
from .forms import CommentForm
from django.core.mail import send_mail
from django.template.loader import render_to_string
from django.http import JsonResponse
from django.conf import settings
from django.urls import reverse

class SendMail(threading.Thread):
    def __init__(self,subject,text,email,fail_silently=False):
        self.subject = subject
        self.text = text
        self.email = email
        self.fail_silently = fail_silently
        threading.Thread.__init__(self)
    def run(self):
        send_mail(
                self.subject,
                '',
                settings.EMAIL_HOST_USER,
                [self.email],
                fail_silently=self.fail_silently,
                html_message=self.text
            )

def update_comment(request):
    referer = request.META.get('HTTP_REFERER', reverse('home'))
    comment_form = CommentForm(request.POST, user =request.user)
    data ={}
    if comment_form.is_valid():
        comment = Comment()
        comment.user = comment_form.cleaned_data['user']
        comment.content = comment_form.cleaned_data['text']
        comment.content_object = comment_form.cleaned_data['content_object']

        parent =  comment_form.cleaned_data['parent']
        if parent:
            comment.parent = parent
            comment.reply = parent.user
            comment.root = parent.root if parent.root else parent
        comment.save()

        #评论发送邮件
        if parent:
            #回复评论
            subject = '有人回复你的评论'
            email = comment.reply.email
        else:
            #发表评论
            subject = '有人评论你的博客'
            email = comment.content_object.get_email()
        content = {}
        content['comment_text'] = comment.content
        content['url'] = comment.content_object.get_url()
        text = render_to_string('send_email.html',content)
        if email != '':
            sendmail = SendMail(subject,text,email)
            sendmail.start()

        data['status'] = 'SUCCESS'
        data['username'] = comment.user.get_nickname_or_username()
        data['comment_time'] = comment.comment_time.strftime('%Y-%m-%d %H:%M:%S')
        data['text'] = comment.content
        if parent:
            data['reply'] = comment.reply.get_nickname_or_username()
        else:
            data['reply'] = ''
        data['pk'] = comment.pk
        data['root_pk'] = comment.root.pk if comment.root else ''
    else:
        data['status'] = 'ERROR'
        data['message'] = list(comment_form.errors.values())[0][0]
    return JsonResponse(data)


__author__ = "Michael Gonzalez"
__email__ = "mmg005@eng.ucsd.edu"

import xml.etree.ElementTree as ETree
import functools
import jinja2
import os
from InoCommenter import *
class ClassGenerator:
    libraries = []
    gadgetron_components = []
    loopBody = []
    name = ""
    funcs = ""
    functionDeclarations = ""

    jinja_env = jinja2.Environment(loader=jinja2.FileSystemLoader( os.path.dirname( __file__ ) ))
    jinja_env.line_comment_prefix = "#//"
    jinja_env.line_statement_prefix = "%"
    variableDeclarations = ""
    jinja_file = "ino_template.jinja.ino"

    def declare_objects(self, objs):
        self.gadgetron_components = objs

    def __init__( self, api, include_str="<>"):
        self.gadgetron_components = []
        self.objInstances = []
        self.setupBody = []
        self.variableDeclarations = []
        self.functionDeclarations = []
        self.loopBody = []
        self.find_name(api)
        self.include_str = include_str
        self.behavior_tree = None

    def appendToSetup(self, lines):
        self.setupBody += lines

    def appendToLoop(self, lines):
        self.loopBody += lines

    def define_functions(self, funcs):
        self.function_definitions = funcs

    def declare_functions(self, funcs):
        self.functionDeclarations = funcs

    def declare_variables(self, vars):
        self.variableDeclarations = vars

    def find_name(self, root):
        for component in root:
            if component.tag == "name": self.name = component.text

    def getClass(self):
        jinja_vars = {
            "gadgetron_objs"        : self.gadgetron_components,
            "variable_declarations" : self.variableDeclarations,
            "function_declarations" : self.functionDeclarations,
            "function_definitions"  : self.function_definitions,
            "setup_body"            : self.setupBody,
            "loop_body"             : self.loopBody,
            "name"                  : self.name,
            "custom_header_file"    : self.name.replace(" ", "-") + '.h',
            "header" : createRobotHeader( self.name ),
            "bt"     : self.behavior_tree,
        }

        template = self.jinja_env.get_template(self.jinja_file)
        code = template.render( jinja_vars )
        return code

if __name__ == "__main__":
    import argparse
    argparser = argparse.ArgumentParser("stuff")
    argparser.add_argument("-g", "--gspec", required=True)
    args = argparser.parse_args()
    api = ETree.parse(args.gspec).getroot()
    generator = ClassGenerator(api)
    generator.getClass()


import random
import pickle
import string
from typing import Union
from collections import Counter
from dataclasses import dataclass, field

@dataclass
class NgrammModel:
    model_filename: str = field(default="model.pkl")
    model_weights: dict = field(default_factory=dict)

    def fit(self, filename: str, n: int) -> tuple[bool, Union[Exception, None]]:
        try:
            with open(filename, 'r', encoding='UTF-8') as reading_file:
                text = reading_file.read()
        except FileNotFoundError as e:
            return False, e
        try:
            text = text.lower().translate(str.maketrans("", "", string.punctuation + "«»—…")).replace('\n\n',
                                                                                '\n').split(" ")
            dict = {}
            for i in range(len(text) - (n - 1)):
                next_words = []
                for j in range(n - 1):
                    next_words.append(text[i + j + 1])
                if text[i] not in dict:
                    dict[text[i]] = [' '.join(next_words)]
                else:
                    dict[text[i]].append(' '.join(next_words))
            for key, val in dict.items():
                for k, v in Counter(val).items():
                    val = (k, v)
                    if val in dict[key]:
                        dict[key] = [val]
                    else:
                        dict[key].append(val)
            for key, val in dict.items():
                filtered_val = [x for x in val if type(x) != str]
                val.clear()
                for x in filtered_val:
                    val.append(x)
            self.model_weights |= dict
            return True, None
        except Exception as e:
            return False, e

    def generate(self, length, prefix=None) -> tuple[bool, Union[Exception, None]]:
        try:
            with open(self.model_filename, 'rb') as f:
                self.model_weights = pickle.load(f)
        except FileNotFoundError as e:
            return False, e
        try:
            if len(self.model_weights) != 0:
                random.seed(random.randint(0, len(self.model_weights)))
                keys = list(self.model_weights.keys())
                variaty_list = []
                word_list = []
                generated_str = ''
                if prefix:
                    prefix = ' '.join(prefix)
                else:
                    prefix = random.choice(list(self.model_weights.keys()))
                generated_str += ' '.join(prefix.split()[:-1])
                if prefix.split()[-1] not in self.model_weights.keys():
                    k = random.choice(keys)
                else:
                    k = prefix.split()[-1]
                    generated_str += " "
                    generated_str += k
                for i in range(length):
                    for j in range(len(self.model_weights[k])):
                        variaty_list.append(self.model_weights[k][j][1])
                        word_list.append(self.model_weights[k][j][0])
                    generated_phrase = random.choices(word_list, weights=variaty_list, k=1)
                    variaty_list.clear()
                    word_list.clear()
                    generated_str += k + " " if i > 0 and k != generated_str.split()[-1] else " "
                    generated_str += generated_phrase[0]
                    k = generated_phrase[0].split()[-1] if generated_phrase[0].split()[
                                                               -1] in self.model_weights else random.choice(keys)
                generated_str = ' '.join(generated_str.split()[:length]).replace("  ", " ")
                generated_str += ' '
            else:
                generated_str = 'Текст слишком короткий для генерации по данной модели'
            yield generated_str
            with open('generated.txt', 'a+') as f:
                f.write(generated_str)
            return True, str
        except Exception as e:
            return False, e

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        with open(file=self.model_filename, mode='wb') as file:
            pickle.dump(self.model_weights, file)


###################################
# File Name : process_pipe.py
###################################
#!/usr/bin/python3

import multiprocessing

def child(pipe):
    p = multiprocessing.current_process()

    msg = "Hello World"
    pipe.send(msg)

    print ("[%s] Send a message to pipe : %s" % (p.name, msg))

def main():
    parent_pipe, child_pipe = multiprocessing.Pipe()

    p = multiprocessing.Process(name="child", target=child, args=(child_pipe,))
    p.start()

    print ("Recieved message : %s" % parent_pipe.recv())
    p.join()

if __name__ == "__main__":
    main()


import os
import tempfile
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, Sequence, TextIO, Tuple, Union

import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

from loguru import logger

from .._base import AbstractAutomation
from ..misc.shell import ShellExecutor
from ..structures import TYPE_PATH, TransferDTO

class RcloneAutomation(AbstractAutomation):
    """
    This is a s3-s3  transfer component using rclone.

    Attributes:
        config: `Path` or `str` or `Dict[str, str]`
            Represents source/dest s3 configuration

        files: `List[str]`
            List of filenames to be transferred.
            (Currently it's not used in RcloneAutomation)

        shell_executor: `misc.shell.ShellExecutor`
            A misc component to execute external shell commands.

        debug: `bool`
            A bool flag to represent if it's a debug mode.

    Note:
        For extra params to the rclone executor,
        we pass them as keyword args through `**params`. Some of the params are:
            - buffer_size (size of buffer)
            - multi_thread_streams (how many chunks is to be created for each
            file?)
            - multi_thread_cutoff (threshold in MB to start chunking a single
            file into multi_thread_streams chunks)
            - ntransfers (number of parallelization for downloads)
            - s3_max_upload_parts (how many chunks at max used to upload to s3?)
            - s3_upload_concurrency (number of parallelization for uploads)
    """

    def __init__(
        self,
        config: Union[Dict[str, str], TYPE_PATH],
        files: Optional[Sequence[TYPE_PATH]] = None,
        shell_executor: Optional[ShellExecutor] = None,
        debug: bool = False,
        **params,
    ) -> None:
        super().__init__(config=config, files=files, debug=debug)

        shell_executor = shell_executor or ShellExecutor()
        assert isinstance(shell_executor, ShellExecutor)
        self.shell_executor = shell_executor

        logger.debug(params)
        self.buffer_size = params.get("buffer_size", 50)
        self.multi_thread_streams = params.get("multi_thread_streams", 10)
        self.multi_thread_cutoff = params.get("multi_thread_cutoff", 50)
        self.ntransfers = params.get("ntransfers", 8)
        self.s3_max_upload_parts = params.get("s3_max_upload_parts", 10)
        self.s3_upload_concurrency = params.get("s3_upload_concurrency", 10)

    def _generate_rclone_cfg(self) -> tempfile.NamedTemporaryFile:
        """
        Generate a temporary config file compatible to rclone.
        """
        source_token = self.config["source_token"]
        source_secret = self.config["source_secret"]
        source_s3_endpoint = self.config["source_s3_endpoint"]

        dest_token = self.config["dest_token"]
        dest_secret = self.config["dest_secret"]
        dest_s3_endpoint = self.config["dest_s3_endpoint"]

        lines = [
            "[s3source]\n",
            "type = s3\n",
            "provider = Other\n",
            f"access_key_id = {source_token}\n",
            f"secret_access_key = {source_secret}\n",
            f"endpoint = {source_s3_endpoint}\n",
            "acl = authenticated-read\n",
            "\n",
            "[s3dest]\n",
            "type = s3\n",
            "provider = Other\n",
            f"access_key_id = {dest_token}\n",
            f"secret_access_key = {dest_secret}\n",
            f"endpoint = {dest_s3_endpoint}\n",
            "acl = authenticated-read\n" "\n",
        ]

        # We don't use context manager here as it will auto-delete the file
        # Or we have to manually delete the file even after closing the
        # file pointer
        ftemp = tempfile.NamedTemporaryFile(mode="w", suffix=".conf", prefix="rclone_")
        ftemp.writelines(lines)
        ftemp.flush()
        return ftemp

    def run_automation(self, **kwargs) -> Tuple[TransferDTO]:
        """
        Main interface to RcloneAutomation.

        Returns:
            tuple of individual file data transfer `Tuple[TransferDTO]`
            where each element object stores
            - filename
            - start_time
            - end_time
            - transferer ("rclone")
        """
        start_automation = time.time()
        source_s3_bucket = self.config["source_s3_bucket"]
        source_s3_region = self.config["source_s3_region"]
        dest_s3_bucket = self.config["dest_s3_bucket"]
        dest_s3_region = self.config["dest_s3_region"]

        # temp files
        rclone_log_file = tempfile.NamedTemporaryFile(
            mode="w+", prefix="rclone_", suffix=".log", delete=False
        )
        logger.debug(f"rclone log file :: {rclone_log_file.name}")

        rclone_config_file = self._generate_rclone_cfg()
        assert rclone_config_file is not None
        logger.debug(f"rclone conf file :: {rclone_config_file.name}")

        cmd = [
            "rclone",
            "copy",
            f"s3source:{source_s3_bucket}",
            f"s3dest:{dest_s3_bucket}",
            f"--multi-thread-streams={self.multi_thread_streams}",
            f"--multi-thread-cutoff={self.multi_thread_cutoff}M",
            f"--s3-max-upload-parts={self.s3_max_upload_parts}",
            f"--s3-upload-concurrency={self.s3_upload_concurrency}",
            f"--buffer-size={self.buffer_size}M",
            f"--transfers={self.ntransfers}",
            "--progress",
            f"--config={rclone_config_file.name}",
            f"--log-file={rclone_log_file.name}",
            "--log-level=DEBUG",
            "-I",
        ]

        start = time.time()
        exdto = self.shell_executor(cmd)
        logger.debug(f"Execution took {time.time()-start} seconds.")

        # this deletes the temp file also
        logger.info(f"Removing temp config at {rclone_config_file.name}")
        rclone_config_file.close()

        # start_time_map, end_time_map = self.parse_log(rclone_log_file, debug=self.debug)
        vals = self.parse_log(rclone_log_file, debug=self.debug)
        logger.debug(
            f"Delta time for {self.__classname__} = {time.time() - start_automation}"
        )
        return vals

    def parse_log(
        self, log: Union[str, TextIO], debug: bool = False
    ) -> Tuple[TransferDTO]:
        """
        Parse rclone-generated log file to extract transfer information
        for each file.

        Returns:
            tuple of data transfer metadata `Tuple[TransferDTO]`.
        """
        # in case it's a path
        if isinstance(log, str):
            log = open(log)
        log.seek(0)

        logger.debug(f"Parsing log at {log.name}")

        timekeeper = {}
        for line in log:
            if (
                line.find("multipart upload starting chunk 1 size") != -1
                or line.find("Transferring unconditionally") != -1
            ):
                fname = line.split(":")[3].strip()
                dto = timekeeper.get(
                    fname, TransferDTO(fname=fname, transferer="rclone")
                )
                dto.start_time = datetime.strptime(
                    line.split("DEBUG")[0].strip(), "%Y/%m/%d %H:%M:%S"
                )
                timekeeper[fname] = dto

            if line.find("Copied") != -1:
                fname = line.split(":")[3].strip()
                dto = timekeeper.get(
                    fname, TransferDTO(fname=fname, transferer="rclone")
                )
                dto.end_time = datetime.strptime(
                    line.split("INFO")[0].strip(), "%Y/%m/%d %H:%M:%S"
                )
                timekeeper[fname] = dto

        if debug:
            logger.debug(f"Transfer maps => {timekeeper}")
        return tuple(timekeeper.values())


import threading
import subprocess
import sqlite3
from datetime import datetime

def count_miners_and_insert_into_db():
    # Run shell commands to count the number of unique miners
    try:
        cmd = "tail -500000 access.log | grep verify | cut -d' ' -f1 | sort | uniq | wc -l"
        total_miners = subprocess.getoutput(cmd)
        total_miners = int(total_miners.strip())  # Convert the output to an integer

        # Connect to the difficulty database
        conn = sqlite3.connect('difficulty.db')
        cursor = conn.cursor()

        # Create the miners table if it doesn't exist
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS miners (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            date DATETIME DEFAULT CURRENT_TIMESTAMP,
            total_miners INTEGER
        );
        ''')

        # Insert the count into the miners table
        cursor.execute("INSERT INTO miners (total_miners) VALUES (?)", (total_miners,))
        conn.commit()

        # Close the database connection
        conn.close()

        print(f"Inserted {total_miners} into the database at {datetime.now()}")
    except Exception as e:
        print(f"An error occurred: {e}")

# Set up a timer to run the function every 60 seconds
def start_timer():
    threading.Timer(60.0, start_timer).start()
    count_miners_and_insert_into_db()

# Start the timer
start_timer()


import time
from utils.py_utils import run_loop
from utils.light_utils import make_multi_grad
from PixelGroup.PixelGroupChain import PixelGroupChain
from consts import REDIS_PORT, LED_CONFIG, LED_COUNT
from collections import namedtuple
import logging
import os
import redis
import json

db = redis.Redis(
	host="localhost",
	port=REDIS_PORT,
	decode_responses=True
)

def from_db():
	chain = PixelGroupChain(LED_CONFIG)

	def read_write_time():
		write_time = float(db.get("lastWrite") or 0)

		if not write_time:
			logging.warn("Could not read lastWrite, using default 0")
			write_time = 0

		return write_time

	def read_colors():
		color_data = json.loads(db.get("solidColor") or "{\"type\":\"solidColor\",\"color\":[0,0,0]}")

		return [color_data["color"]]

	def render_pixels(colors):
		chain.fade_to_all(make_multi_grad(colors, LED_COUNT), 0.5)

	render_pixels(read_colors())
	last_write_time = read_write_time()

	def loop():
		nonlocal last_write_time

		write_time = read_write_time()

		if write_time > last_write_time:
			render_pixels(read_colors())
			last_write_time = write_time

		time.sleep(1)

	run_loop(loop)


from rest_framework import serializers
from .models import Tag, Good, Category

class TagSerializer(serializers.ModelSerializer):
    name = serializers.JSONField()

    class Meta:
        model = Tag
        fields = ['id', 'name', 'uuid']

class GoodSerializer(serializers.ModelSerializer):

    class Meta:
        model = Good
        fields = ['id', 'name', 'description', 'price', 'active', 'country', 'category', 'tags', 'image']

class CategorySerializer(serializers.ModelSerializer):
    goods = GoodSerializer(many=True, read_only=True)
    good_count = serializers.SerializerMethodField()

    class Meta:
        model = Category
        fields = ['id', 'name', 'description', 'goods', 'good_count']

    def get_good_count(self, obj):
        count = obj.goods.count()
        return count


'''Parser.

Notes:
    Main executable script.

Run as:

    $ python3 parse.py --paper
    $ python3 parse.py --aer1217
    $ python3 parse.py

'''
import os
import argparse
import csv
import sys

import matplotlib.pyplot as plt
import numpy as np
import pybullet as p
import tikzplotlib
from scipy.interpolate import interp1d

# import bagpy
# from bagpy import bagreader
# import pickle

def run(
    paper=False,
    aer1217=False,
):
    '''Main function.
    '''
    sol_num = -1
    if paper:
        fig, axs = plt.subplots(3, 9)
    for sol in ['arg', 'eku', 'h2']:
        sol_num += 1
        if not paper and not aer1217:
            fig, axs = plt.subplots(16, 6)
        num_files = 10

        real_ref_time = [[] for i in range(num_files)]
        real_ref_x = [[] for i in range(num_files)]
        real_ref_y = [[] for i in range(num_files)]
        real_ref_z = [[] for i in range(num_files)]
        real_time = [[] for i in range(num_files)]
        real_x = [[] for i in range(num_files)]
        real_y = [[] for i in range(num_files)]
        real_z = [[] for i in range(num_files)]
        real_q1 = [[] for i in range(num_files)]
        real_q2 = [[] for i in range(num_files)]
        real_q3 = [[] for i in range(num_files)]
        real_q4 = [[] for i in range(num_files)]

        real_ref_vx = [[] for i in range(num_files)]
        real_ref_vy = [[] for i in range(num_files)]
        real_ref_vz = [[] for i in range(num_files)]
        real_ref_ax = [[] for i in range(num_files)]
        real_ref_ay = [[] for i in range(num_files)]
        real_ref_az = [[] for i in range(num_files)]
        real_ref_yaw = [[] for i in range(num_files)]
        real_ref_w1 = [[] for i in range(num_files)]
        real_ref_w2 = [[] for i in range(num_files)]
        real_ref_w3 = [[] for i in range(num_files)]

        sim_time = [[] for i in range(num_files)]
        sim_ref_x = [[] for i in range(num_files)]
        sim_ref_y = [[] for i in range(num_files)]
        sim_ref_z = [[] for i in range(num_files)]
        sim_x = [[] for i in range(num_files)]
        sim_y = [[] for i in range(num_files)]
        sim_z = [[] for i in range(num_files)]
        sim_r = [[] for i in range(num_files)]
        sim_p = [[] for i in range(num_files)]
        sim_j = [[] for i in range(num_files)]

        sim_ref_vx = [[] for i in range(num_files)]
        sim_ref_vy = [[] for i in range(num_files)]
        sim_ref_vz = [[] for i in range(num_files)]
        sim_ref_ax = [[] for i in range(num_files)]
        sim_ref_ay = [[] for i in range(num_files)]
        sim_ref_az = [[] for i in range(num_files)]
        sim_ref_yaw = [[] for i in range(num_files)]
        sim_ref_w1 = [[] for i in range(num_files)]
        sim_ref_w2 = [[] for i in range(num_files)]
        sim_ref_w3 = [[] for i in range(num_files)]

        avg_real_cmd_timestep = 0

        t_new = [[] for i in range(num_files)]
        initial_skip = [0 for i in range(num_files)]
        resampled_real_ref_x = [[] for i in range(num_files)]
        resampled_real_ref_y = [[] for i in range(num_files)]
        resampled_real_ref_z = [[] for i in range(num_files)]
        resampled_real_x = [[] for i in range(num_files)]
        resampled_real_y = [[] for i in range(num_files)]
        resampled_real_z = [[] for i in range(num_files)]
        resampled_real_r = [[] for i in range(num_files)]
        resampled_real_p = [[] for i in range(num_files)]
        resampled_real_j = [[] for i in range(num_files)]

        se_x = 0
        se_y = 0
        se_z = 0
        se_count = 0
        se_ref_x = 0
        se_ref_y = 0
        se_ref_z = 0
        se_ref_count = 0

        for num in range(1, num_files+1):
            idx = num-1

            ############################
            ############################
            ############################

            # real cmd
            first_cmd_time = None
            last_cmd_time = None
            with open('./exp/cmd_test_' + sol + str(num) + '.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    if i == 0:
                        continue
                    if i == 1:
                        real_ref_init_time = float(row[0])
                        first_cmd_time = real_ref_init_time
                    real_ref_time[idx].append(float(row[0])-real_ref_init_time)
                    if row[1] == 'takeoff' or row[1] == 'none':
                        pass
                    real_ref_x[idx].append(float(row[2]))
                    real_ref_y[idx].append(float(row[3]))
                    real_ref_z[idx].append(float(row[4]))
                    #
                    real_ref_vx[idx].append(float(row[5]))
                    real_ref_vy[idx].append(float(row[6]))
                    real_ref_vz[idx].append(float(row[7]))
                    real_ref_ax[idx].append(float(row[8]))
                    real_ref_ay[idx].append(float(row[9]))
                    real_ref_az[idx].append(float(row[10]))
                    real_ref_yaw[idx].append(float(row[11]))
                    real_ref_w1[idx].append(float(row[12]))
                    real_ref_w2[idx].append(float(row[13]))
                    real_ref_w3[idx].append(float(row[14]))
                last_cmd_time = real_ref_time[idx][-1]+real_ref_init_time

            # real vicon
            with open('./exp/' + sol + str(num) + '.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    if i == 0:
                        continue
                    time = float(row[0])
                    if time >= first_cmd_time and time <= last_cmd_time:
                        real_time[idx].append(time-real_ref_init_time)
                        real_x[idx].append(float(row[5]))
                        real_y[idx].append(float(row[6]))
                        real_z[idx].append(float(row[7]))
                        real_q1[idx].append(float(row[8]))
                        real_q2[idx].append(float(row[9]))
                        real_q3[idx].append(float(row[10]))
                        real_q4[idx].append(float(row[11]))

            # sim cmd
            sim_ref_time = []
            with open('./sim/' + sol + '/episode' + str(num) + '.csv',
                      encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    if i == 0:
                        continue
                    sim_ref_time.append(float(row[0]))
                    if row[1] == 'takeoff' or row[1] == 'none':
                        pass
                    sim_ref_x[idx].append(float(row[2]))
                    sim_ref_y[idx].append(float(row[3]))
                    sim_ref_z[idx].append(float(row[4]))
                    #
                    sim_ref_vx[idx].append(float(row[5]))
                    sim_ref_vy[idx].append(float(row[6]))
                    sim_ref_vz[idx].append(float(row[7]))
                    sim_ref_ax[idx].append(float(row[8]))
                    sim_ref_ay[idx].append(float(row[9]))
                    sim_ref_az[idx].append(float(row[10]))
                    sim_ref_yaw[idx].append(float(row[11]))
                    sim_ref_w1[idx].append(float(row[12]))
                    sim_ref_w2[idx].append(float(row[13]))
                    sim_ref_w3[idx].append(float(row[14]))

            # sim pos
            folder = './sim/' + sol + '/ep' + str(num)
            sim_x_time = []
            sim_y_time = []
            sim_z_time = []
            with open(folder + '/x0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_x_time.append(float(row[0]))
                    sim_x[idx].append(float(row[1]))
            with open(folder + '/y0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_y_time.append(float(row[0]))
                    sim_y[idx].append(float(row[1]))
            with open(folder + '/z0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_z_time.append(float(row[0]))
                    sim_z[idx].append(float(row[1]))
            if sim_x_time != sim_y_time or sim_x_time != sim_z_time or sim_x_time != sim_ref_time:
                print('Error: time mismatch')
                exit()
            sim_time[idx] = sim_x_time

            with open(folder + '/r0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_r[idx].append(float(row[1]))
            with open(folder + '/p0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_p[idx].append(float(row[1]))
            with open(folder + '/ya0.csv', encoding='utf-8') as csv_file:
                csv_reader = csv.reader(csv_file, delimiter=',')
                for i, row in enumerate(csv_reader):
                    sim_j[idx].append(float(row[1]))

            ############################
            ############################
            ############################

            # math magic
            initial_skip[idx] = 0
            for val in sim_time[idx]:
                if val >= real_ref_time[idx][0] and val >= real_time[idx][0]:
                    t_new[idx].append(val)
                else:
                    initial_skip[idx] += 1
                if val >= real_ref_time[idx][-1] or val >= real_time[idx][-1]:
                    print('Something is not right')
                    sys.exit()
            t_new[idx] = np.array(t_new[idx])

            real_ref_x_func = interp1d(
                np.array(real_ref_time[idx]), np.array(real_ref_x[idx]), kind='linear')
            resampled_real_ref_x[idx] = real_ref_x_func(t_new[idx])
            real_ref_y_func = interp1d(
                np.array(real_ref_time[idx]), np.array(real_ref_y[idx]), kind='linear')
            resampled_real_ref_y[idx] = real_ref_y_func(t_new[idx])
            real_ref_z_func = interp1d(
                np.array(real_ref_time[idx]), np.array(real_ref_z[idx]), kind='linear')
            resampled_real_ref_z[idx] = real_ref_z_func(t_new[idx])

            real_x_func = interp1d(
                np.array(real_time[idx]), np.array(real_x[idx]), kind='linear')
            resampled_real_x[idx] = real_x_func(t_new[idx])
            real_y_func = interp1d(
                np.array(real_time[idx]), np.array(real_y[idx]), kind='linear')
            resampled_real_y[idx] = real_y_func(t_new[idx])
            real_z_func = interp1d(
                np.array(real_time[idx]), np.array(real_z[idx]), kind='linear')
            resampled_real_z[idx] = real_z_func(t_new[idx])

            ############################
            ############################
            ############################

            # plot one-by-one
            if not paper and not aer1217:
                axs[0, 0].plot(real_ref_time[idx],
                               real_ref_x[idx], label='ref')
                axs[1, 0].plot(real_ref_time[idx],
                               real_ref_y[idx], label='ref')
                axs[2, 0].plot(real_ref_time[idx],
                               real_ref_z[idx], label='ref')
            #
                axs[0, 0].plot(
                    t_new[idx], resampled_real_ref_x[idx], label='ref')
                axs[1, 0].plot(
                    t_new[idx], resampled_real_ref_y[idx], label='ref')
                axs[2, 0].plot(
                    t_new[idx], resampled_real_ref_z[idx], label='ref')
            #
                axs[6, 0].plot(real_ref_time[idx],
                               real_ref_vx[idx], label='ref')
                axs[7, 0].plot(real_ref_time[idx],
                               real_ref_vy[idx], label='ref')
                axs[8, 0].plot(real_ref_time[idx],
                               real_ref_vz[idx], label='ref')
                axs[9, 0].plot(real_ref_time[idx],
                               real_ref_ax[idx], label='ref')
                axs[10, 0].plot(real_ref_time[idx],
                                real_ref_ay[idx], label='ref')
                axs[11, 0].plot(real_ref_time[idx],
                                real_ref_az[idx], label='ref')
                axs[12, 0].plot(real_ref_time[idx],
                                real_ref_yaw[idx], label='ref')
                axs[13, 0].plot(real_ref_time[idx],
                                real_ref_w1[idx], label='ref')
                axs[14, 0].plot(real_ref_time[idx],
                                real_ref_w2[idx], label='ref')
                axs[15, 0].plot(real_ref_time[idx],
                                real_ref_w3[idx], label='ref')

            avg_time = (
                real_ref_time[idx][-1] - real_ref_time[idx][0]) / (len(real_ref_time[idx])-1)
            avg_real_cmd_timestep += avg_time

            if paper:
                axs[0, 0+3*sol_num].plot(real_time[idx],
                                         real_x[idx], label='exp')
                axs[1, 0+3*sol_num].plot(real_time[idx],
                                         real_y[idx], label='exp')
                axs[2, 0+3*sol_num].plot(real_time[idx],
                                         real_z[idx], label='exp')
            elif not aer1217:
                axs[0, 1].plot(real_time[idx], real_x[idx], label='exp')
                axs[1, 1].plot(real_time[idx], real_y[idx], label='exp')
                axs[2, 1].plot(real_time[idx], real_z[idx], label='exp')
                #
                axs[0, 1].plot(t_new[idx], resampled_real_x[idx], label='exp')
                axs[1, 1].plot(t_new[idx], resampled_real_y[idx], label='exp')
                axs[2, 1].plot(t_new[idx], resampled_real_z[idx], label='exp')
            #
            real_euler1 = []
            real_euler2 = []
            real_euler3 = []
            for time in range(len(real_q1[idx])):
                val = p.getEulerFromQuaternion([
                        real_q1[idx][time],
                        real_q2[idx][time],
                        real_q3[idx][time],
                        real_q4[idx][time]
                     ])
                real_euler1.append(val[0])
                real_euler2.append(val[1])
                real_euler3.append(val[2])
                # getEulerFromQuaternion returns a list of 3 floating point values, a vec3.
                # The rotation order is first roll around X,
                # then pitch around Y and finally yaw around Z,
                # as in the ROS URDF rpy convention.
            if not paper and not aer1217:
                axs[3, 1].plot(real_time[idx], real_euler1, label='exp')
                axs[4, 1].plot(real_time[idx], real_euler2, label='exp')
                axs[5, 1].plot(real_time[idx], real_euler3, label='exp')
            #
            real_r_func = interp1d(
                np.array(real_time[idx]), np.array(real_euler1), kind='linear')
            resampled_real_r[idx] = real_r_func(t_new[idx])
            real_p_func = interp1d(
                np.array(real_time[idx]), np.array(real_euler2), kind='linear')
            resampled_real_p[idx] = real_p_func(t_new[idx])
            real_j_func = interp1d(
                np.array(real_time[idx]), np.array(real_euler3), kind='linear')
            resampled_real_j[idx] = real_j_func(t_new[idx])
            #
            if not paper and not aer1217:
                axs[3, 1].plot(t_new[idx], resampled_real_r[idx], label='exp')
                axs[4, 1].plot(t_new[idx], resampled_real_p[idx], label='exp')
                axs[5, 1].plot(t_new[idx], resampled_real_j[idx], label='exp')

                axs[0, 2].plot(sim_time[idx], sim_ref_x[idx], label='ref')
                axs[1, 2].plot(sim_time[idx], sim_ref_y[idx], label='ref')
                axs[2, 2].plot(sim_time[idx], sim_ref_z[idx], label='ref')

                axs[6, 2].plot(sim_time[idx], sim_ref_vx[idx], label='ref')
                axs[7, 2].plot(sim_time[idx], sim_ref_vy[idx], label='ref')
                axs[8, 2].plot(sim_time[idx], sim_ref_vz[idx], label='ref')
                axs[9, 2].plot(sim_time[idx], sim_ref_ax[idx], label='ref')
                axs[10, 2].plot(sim_time[idx], sim_ref_ay[idx], label='ref')
                axs[11, 2].plot(sim_time[idx], sim_ref_az[idx], label='ref')
                axs[12, 2].plot(sim_time[idx], sim_ref_yaw[idx], label='ref')
                axs[13, 2].plot(sim_time[idx], sim_ref_w1[idx], label='ref')
                axs[14, 2].plot(sim_time[idx], sim_ref_w2[idx], label='ref')
                axs[15, 2].plot(sim_time[idx], sim_ref_w3[idx], label='ref')

            if paper:
                axs[0, 1+3*sol_num].plot(sim_time[idx],
                                         sim_x[idx], label='sim')
                axs[1, 1+3*sol_num].plot(sim_time[idx],
                                         sim_y[idx], label='sim')
                axs[2, 1+3*sol_num].plot(sim_time[idx],
                                         sim_z[idx], label='sim')
            elif not aer1217:
                axs[0, 3].plot(sim_time[idx], sim_x[idx], label='sim')
                axs[1, 3].plot(sim_time[idx], sim_y[idx], label='sim')
                axs[2, 3].plot(sim_time[idx], sim_z[idx], label='sim')
                #
                axs[3, 3].plot(sim_time[idx], sim_r[idx], label='sim')
                axs[4, 3].plot(sim_time[idx], sim_p[idx], label='sim')
                axs[5, 3].plot(sim_time[idx], sim_j[idx], label='sim')

            se_x += np.sum((resampled_real_x[idx] -
                           np.array(sim_x[idx][initial_skip[idx]:]))**2)
            se_y += np.sum((resampled_real_y[idx] -
                           np.array(sim_y[idx][initial_skip[idx]:]))**2)
            se_z += np.sum((resampled_real_z[idx] -
                           np.array(sim_z[idx][initial_skip[idx]:]))**2)
            se_count += len(resampled_real_x[idx])

            se_ref_x += np.sum((resampled_real_ref_x[idx] -
                               np.array(sim_ref_x[idx][initial_skip[idx]:]))**2)
            se_ref_y += np.sum((resampled_real_ref_y[idx] -
                               np.array(sim_ref_y[idx][initial_skip[idx]:]))**2)
            se_ref_z += np.sum((resampled_real_ref_z[idx] -
                               np.array(sim_ref_z[idx][initial_skip[idx]:]))**2)
            se_ref_count += len(resampled_real_ref_x[idx])

            if not paper and not aer1217:
                axs[0, 4].plot(
                    t_new[idx], (sim_x[idx][initial_skip[idx]:]-resampled_real_x[idx])**2,
                    label='diff')
                axs[1, 4].plot(
                    t_new[idx], (sim_y[idx][initial_skip[idx]:]-resampled_real_y[idx])**2,
                    label='diff')
                axs[2, 4].plot(
                    t_new[idx], (sim_z[idx][initial_skip[idx]:]-resampled_real_z[idx])**2,
                    label='diff')

                axs[3, 4].plot(
                    t_new[idx], (sim_r[idx][initial_skip[idx]:]-resampled_real_r[idx])**2,
                    label='diff')
                axs[4, 4].plot(
                    t_new[idx], (sim_p[idx][initial_skip[idx]:]-resampled_real_p[idx])**2,
                    label='diff')
                axs[5, 4].plot(
                    t_new[idx], (sim_j[idx][initial_skip[idx]:]-resampled_real_j[idx])**2,
                    label='diff')

                axs[0, 5].plot(
                    t_new[idx], (sim_ref_x[idx][initial_skip[idx]:]-resampled_real_ref_x[idx])**2,
                    label='diff')
                axs[1, 5].plot(
                    t_new[idx], (sim_ref_y[idx][initial_skip[idx]:]-resampled_real_ref_y[idx])**2,
                    label='diff')
                axs[2, 5].plot(
                    t_new[idx], (sim_ref_z[idx][initial_skip[idx]:]-resampled_real_ref_z[idx])**2,
                    label='diff')

            ############################
            ############################
            ############################

        latest_start = -1
        earliest_end = 9000
        for i in range(num_files):
            if sol == 'eku' and i == 1:
                continue
            latest_start = max(latest_start, t_new[i][0])
            earliest_end = min(earliest_end, t_new[i][-1])
            # print(t_new[i][0], t_new[i][-1])
        print('start and end time for rmse', latest_start, earliest_end)
        for i in range(num_files):
            pass
            # print(f'{initial_skip[i]} {t_new[i][0]:.3f}', t_new[i].shape, np.array(sim_x[i][initial_skip[i]:]).shape, np.array(sim_ref_x[i][initial_skip[i]:]).shape, np.array(resampled_real_x[i]).shape, np.array(resampled_real_ref_x[i]).shape)

        cropped_real_x = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_x = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_real_y = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_y = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_real_z = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_z = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        #
        cropped_real_r = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_r = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_real_p = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_p = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_real_j = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))
        cropped_sim_j = np.zeros(
            (num_files, int((earliest_end-latest_start)*30)))

        for i in range(num_files):
            if sol == 'eku' and i == 1:
                continue
            insert_idx = 0
            for j in range(len(t_new[i])):
                if t_new[i][j] >= latest_start:
                    cropped_real_x[i][insert_idx] = resampled_real_x[i][j]
                    cropped_sim_x[i][insert_idx] = sim_x[i][j]
                    cropped_real_y[i][insert_idx] = resampled_real_y[i][j]
                    cropped_sim_y[i][insert_idx] = sim_y[i][j]
                    cropped_real_z[i][insert_idx] = resampled_real_z[i][j]
                    cropped_sim_z[i][insert_idx] = sim_z[i][j]
                    #
                    cropped_real_r[i][insert_idx] = resampled_real_r[i][j]
                    cropped_sim_r[i][insert_idx] = sim_r[i][j]
                    cropped_real_p[i][insert_idx] = resampled_real_p[i][j]
                    cropped_sim_p[i][insert_idx] = sim_p[i][j]
                    cropped_real_j[i][insert_idx] = resampled_real_j[i][j]
                    cropped_sim_j[i][insert_idx] = sim_j[i][j]
                    insert_idx += 1
                    if insert_idx == cropped_real_x.shape[1]:
                        break

        avg_real_x = np.average(cropped_real_x, axis=0)
        avg_sim_x = np.average(cropped_sim_x, axis=0)
        avg_rmse_x = np.sqrt(np.sum((avg_real_x-avg_sim_x)**2)/len(avg_real_x))
        avg_real_y = np.average(cropped_real_y, axis=0)
        avg_sim_y = np.average(cropped_sim_y, axis=0)
        avg_rmse_y = np.sqrt(np.sum((avg_real_y-avg_sim_y)**2)/len(avg_real_y))
        avg_real_z = np.average(cropped_real_z, axis=0)
        avg_sim_z = np.average(cropped_sim_z, axis=0)
        avg_rmse_z = np.sqrt(np.sum((avg_real_z-avg_sim_z)**2)/len(avg_real_z))
        #
        avg_real_r = np.average(cropped_real_r, axis=0)
        avg_sim_r = np.average(cropped_sim_r, axis=0)
        avg_rmse_r = np.sqrt(np.sum((avg_real_r-avg_sim_r)**2)/len(avg_real_r))
        avg_real_p = np.average(cropped_real_p, axis=0)
        avg_sim_p = np.average(cropped_sim_p, axis=0)
        avg_rmse_p = np.sqrt(np.sum((avg_real_p-avg_sim_p)**2)/len(avg_real_p))
        avg_real_j = np.average(cropped_real_j, axis=0)
        avg_sim_j = np.average(cropped_sim_j, axis=0)
        avg_rmse_j = np.sqrt(np.sum((avg_real_j-avg_sim_j)**2)/len(avg_real_j))
        print(
            f'avg_rmse_xyz {avg_rmse_x:.2f} {avg_rmse_y:.2f} {avg_rmse_z:.2f}')
        print(
            f'avg_rmse_rpj {avg_rmse_r:.2f} {avg_rmse_p:.2f} {avg_rmse_j:.2f}')

        cropped_time = np.linspace(
            latest_start, earliest_end, cropped_real_x.shape[1])
        if paper:
            axs[0, 2+3*sol_num].plot(cropped_time,
                                     np.sqrt((avg_real_x-avg_sim_x)**2), label='diff')
            axs[1, 2+3*sol_num].plot(cropped_time,
                                     np.sqrt((avg_real_y-avg_sim_y)**2), label='diff')
            axs[2, 2+3*sol_num].plot(cropped_time,
                                     np.sqrt((avg_real_z-avg_sim_z)**2), label='diff')
            #
            axs[0, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_x-avg_sim_x)**2)
                                     +np.std(np.sqrt((cropped_real_x-cropped_sim_x)**2), axis=0),
                                     label='diff')
            axs[0, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_x-avg_sim_x)**2)
                                     -np.std(np.sqrt((cropped_real_x-cropped_sim_x)**2), axis=0),
                                     label='diff')
            axs[1, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_y-avg_sim_y)**2)
                                     +np.std(np.sqrt((cropped_real_y-cropped_sim_y)**2), axis=0),
                                     label='diff')
            axs[1, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_y-avg_sim_y)**2)
                                     -np.std(np.sqrt((cropped_real_y-cropped_sim_y)**2), axis=0),
                                     label='diff')
            axs[2, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_z-avg_sim_z)**2)
                                     +np.std(np.sqrt((cropped_real_z-cropped_sim_z)**2), axis=0),
                                     label='diff')
            axs[2, 2+3*sol_num].plot(cropped_time, np.sqrt((avg_real_z-avg_sim_z)**2)
                                     -np.std(np.sqrt((cropped_real_z-cropped_sim_z)**2), axis=0),
                                     label='diff')
        elif not aer1217:
            axs[0, 4].plot(cropped_time, (avg_real_x-avg_sim_x)
                           ** 2, label='diff')
            axs[1, 4].plot(cropped_time, (avg_real_y-avg_sim_y)
                           ** 2, label='diff')
            axs[2, 4].plot(cropped_time, (avg_real_z-avg_sim_z)
                           ** 2, label='diff')
            #
            axs[3, 4].plot(cropped_time, (avg_real_r-avg_sim_r)
                           ** 2, label='diff')
            axs[4, 4].plot(cropped_time, (avg_real_p-avg_sim_p)
                           ** 2, label='diff')
            axs[5, 4].plot(cropped_time, (avg_real_j-avg_sim_j)
                           ** 2, label='diff')

        # mse_ref_x = se_ref_x/se_ref_count
        # rmse_ref_x = np.sqrt(mse_ref_x)
        # mse_ref_y = se_ref_y/se_ref_count
        # rmse_ref_y = np.sqrt(mse_ref_y)
        # mse_ref_z = se_ref_z/se_ref_count
        # rmse_ref_z = np.sqrt(mse_ref_z)
        # avg_real_cmd_timestep /= 10
        # real_cmd_freq =  1/avg_real_cmd_timestep
        # print(f'rmse_ref_xyz ({rmse_ref_x:.2f} {rmse_ref_y:.2f} {rmse_ref_z:.2f}) real_cmd_freq {real_cmd_freq:.2f}')

        # rmse_x = np.sqrt(se_x/se_count)
        # rmse_y = np.sqrt(se_y/se_count)
        # rmse_z = np.sqrt(se_z/se_count)

        # labels
        # fig.suptitle(f'{sol}, rmse_xyz ({rmse_x:.2f} {rmse_y:.2f} {rmse_z:.2f})')
        if not paper and not aer1217:
            fig.suptitle(
                f'{sol}, avg rmse xyz {avg_rmse_x:.2f} {avg_rmse_y:.2f} {avg_rmse_z:.2f} \
                rpj  {avg_rmse_r:.2f} {avg_rmse_p:.2f} {avg_rmse_j:.2f}')
            axs[5, 0].set_xlabel('time')
            axs[5, 1].set_xlabel('time')
            axs[5, 2].set_xlabel('time')
            axs[5, 3].set_xlabel('time')
            axs[5, 4].set_xlabel('time')
            axs[5, 5].set_xlabel('time')

            axs[0, 0].set_ylabel('x')
            axs[1, 0].set_ylabel('y')
            axs[2, 0].set_ylabel('z')
            axs[3, 0].set_ylabel('r')
            axs[4, 0].set_ylabel('p')
            axs[5, 0].set_ylabel('j')

            axs[6, 0].set_ylabel('vx')
            axs[7, 0].set_ylabel('vy')
            axs[8, 0].set_ylabel('vz')
            axs[9, 0].set_ylabel('ax')
            axs[10, 0].set_ylabel('ay')
            axs[11, 0].set_ylabel('az')
            axs[12, 0].set_ylabel('yaw')
            axs[13, 0].set_ylabel('w1')
            axs[14, 0].set_ylabel('w2')
            axs[15, 0].set_ylabel('w3')

            axs[0, 0].set_title('real_cmd')
            axs[0, 1].set_title('real')
            axs[0, 2].set_title('sim_cmd')
            axs[0, 3].set_title('sim')
            axs[0, 4].set_title('sq_diff')
            axs[0, 5].set_title('cmd_sq_diff')

            # axes limits
            y_axis_range = 7
            axs[0, 0].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[0, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[0, 2].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[0, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[0, 4].set_ylim(0, 1.25*y_axis_range)
            axs[0, 5].set_ylim(0, 1.25*y_axis_range)

            axs[1, 0].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[1, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[1, 2].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[1, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[1, 4].set_ylim(0, 1.25*y_axis_range)
            axs[1, 5].set_ylim(0, 1.25*y_axis_range)

            axs[2, 0].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[2, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[2, 2].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[2, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
            axs[2, 4].set_ylim(0, 1.25*y_axis_range)
            axs[2, 5].set_ylim(0, 1.25*y_axis_range)

            if sol == 'arg':
                axs[3, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 4].set_ylim(0, 1.25*y_axis_range)
                axs[4, 4].set_ylim(0, 1.25*y_axis_range)
                axs[5, 4].set_ylim(0, 3*y_axis_range)
            elif sol == 'eku':
                axs[3, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 4].set_ylim(0, 1.25*y_axis_range)
                axs[4, 4].set_ylim(0, 1.25*y_axis_range)
                axs[5, 4].set_ylim(0, 1.25*y_axis_range)
            if sol == 'h2':
                axs[3, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 1].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[4, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[5, 3].set_ylim(-y_axis_range/2, y_axis_range/2)
                axs[3, 4].set_ylim(0, 1.25*y_axis_range)
                axs[4, 4].set_ylim(0, 1.25*y_axis_range)
                axs[5, 4].set_ylim(0, 5.5*y_axis_range)

        if not paper and not aer1217:
            tikzplotlib.clean_figure(target_resolution=10, scale_precision=1.0)
            tikzplotlib.save('./tikz/' + sol + '.tex')
            plt.savefig('./png/' + sol + '.png')
            plt.show()

    if paper:
        axs[0, 0].set_ylabel('x (m)')
        axs[1, 0].set_ylabel('y (m)')
        axs[2, 0].set_ylabel('z (m)')
        for col in range(9):
            axs[2, col].set_xlabel('time (s)')
        for col in range(1, 9):
            for row in range(3):
                axs[row, col].yaxis.set_ticklabels([])
        for col in range(9):
            for row in range(2):
                axs[row, col].xaxis.set_ticklabels([])
        axs[0, 0].set_title('real arg')
        axs[0, 3].set_title('real eku')
        axs[0, 6].set_title('real h2')
        axs[0, 1].set_title('sim arg')
        axs[0, 4].set_title('sim eku')
        axs[0, 7].set_title('sim h2')
        for col in [2, 5, 8]:
            axs[0, col].set_title('diff (avg, pm std)')
        for col in [0, 1, 3, 4, 6, 7]:
            for row in range(3):
                axs[row, col].set_ylim(-3.5, 3.5)
        for col in [2, 5, 8]:
            for row in range(3):
                axs[row, col].set_ylim(-2, 6)
        for col in [0, 1, 2]:
            for row in range(3):
                axs[row, col].set_xlim(0, 12)
        for col in [3, 4, 5]:
            for row in range(3):
                axs[row, col].set_xlim(0, 24)
        for col in [6, 7, 8]:
            for row in range(3):
                axs[row, col].set_xlim(0, 12)

        tikzplotlib.clean_figure(target_resolution=300, scale_precision=1.0)
        tikzplotlib.save('./tikz/all.tex')
        plt.savefig('./png/all.png')
        plt.show()

    if aer1217:
        fig = plt.figure()
        axs = []
        for i in range(5):
            axs.append(fig.add_subplot(151+i, projection='3d'))
            # axs[i] = plt.axes(projection='3d')
            # axs[i].view_init(elev=20., azim=-35, roll=0)
            axs[i].set_title(str(i))
        for root, _, files in os.walk('./aer1217/data/baseline/'):
            for file in files:
                if file.endswith('.csv'):
                    x, y, z = [], [], []
                    with open(os.path.join(root, file), encoding='utf-8') as csv_file:
                        csv_reader = csv.reader(csv_file, delimiter=',')
                        for i, row in enumerate(csv_reader):
                            if i == 0:
                                continue
                            x.append(float(row[2]))
                            y.append(float(row[3]))
                            z.append(float(row[4]))
                    if 'A' in file:
                        axs[0].plot(x, y, z, label=file)
                    if 'C' in file:
                        axs[1].plot(x, y, z, label=file)
                    if 'D' in file and '6' not in file:
                        axs[2].plot(x, y, z, label=file)
                    if 'F' in file:
                        axs[3].plot(x, y, z, label=file)
                    if 'G' in file:
                        axs[4].plot(x, y, z, label=file)
        tikzplotlib.clean_figure(target_resolution=300, scale_precision=1.0)
        tikzplotlib.save('./tikz/aer1217.tex')
        plt.savefig('./png/aer1217.png')
        plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Parsing script')
    parser.add_argument('--paper',
                        action='store_true',
                        help='Test argument',
                        )
    parser.add_argument('--aer1217',
                        action='store_true',
                        help='Test argument',
                        )
    ARGS = parser.parse_args()
    run(**vars(ARGS))