#Credit to WeBots for the Mavic controller in C for robot movement: https://github.com/cyberbotics/webots/blob/released/projects/robots/dji/mavic/controllers/mavic2pro/mavic2pro.c
import math
import time
import random
import copy
import cv2
import numpy as np
from matplotlib import pyplot as plt
from controller import Robot, Motor, CameraRecognitionObject, DistanceSensor, Camera, InertialUnit, GPS, Compass, Gyro, Keyboard, Lidar, Display, RangeFinder, LED
LIDAR_SENSOR_MAX_RANGE = 3 # Meters
LIDAR_ANGLE_BINS = 21 # 21 Bins to cover the angular range of the lidar, centered at 10
LIDAR_ANGLE_RANGE = 1.5708 # 90 degrees, 1.5708 radians
pose_x = 0.197
pose_y = 0.678
pose_theta = 0
robot=Robot()
timestep = int(robot.getBasicTimeStep())

camera = robot.getDevice("camera");
Camera.enable(camera, timestep)
Camera.recognitionEnable(camera,timestep)
width = camera.getWidth()
height = camera.getHeight()
def CLAMP(value, low, high):
    if value < low:
        return (low)
    elif (value) > (high):
        return high
    else:
        return (value)

#camera.enable(camera, timestep);
front_left_led = robot.getDevice("front left led");
front_right_led = robot.getDevice("front right led");
imu = robot.getDevice("inertial unit");
gps = robot.getDevice("gps")
InertialUnit.enable(imu, timestep)
GPS.enable(gps, timestep)
compass = robot.getDevice("compass")
Compass.enable(compass, timestep)
gyro = robot.getDevice("gyro")
Gyro.enable(gyro, timestep);
keyboard =Keyboard()
Keyboard.enable(keyboard, timestep);
display=robot.getDevice('display')
camera_roll_motor = robot.getDevice("camera roll");
camera_pitch_motor = robot.getDevice("camera pitch");
front_left_motor = robot.getDevice("front left propeller");
front_right_motor = robot.getDevice("front right propeller");
rear_left_motor = robot.getDevice("rear left propeller");
rear_right_motor = robot.getDevice("rear right propeller");
motors = [front_left_motor, front_right_motor, rear_left_motor, rear_right_motor];
for m in range(4):
    Motor.setPosition(motors[m], float('inf'));
    Motor.setVelocity(motors[m], 1.0);
print("Start the drone...\n");
print("You can control the drone with your computer keyboard:\n");
print("- 'up': move forward.\n");
print("- 'down': move backward.\n");
print("- 'right': turn right.\n");
print("- 'left': turn left.\n");
print("- 'shift + up': increase the target altitude.\n");
print("- 'shift + down': decrease the target altitude.\n");
print("- 'shift + right': strafe right.\n");
print("- 'shift + left': strafe left.\n");

k_vertical_thrust = 68.5;  # with this thrust, the drone lifts.
k_vertical_offset = 0.6;   # Vertical offset where the robot actually targets to stabilize itself.
k_vertical_p = 3.0;        # P constant of the vertical PID.
k_roll_p = 50.0;           # P constant of the roll PID.
k_pitch_p = 30.0;          # P constant of the pitch PID.

target_altitude = 5;  # The target altitude. Can be changed by the user.
cv2.startWindowThread()
cv2.namedWindow("preview")
while (robot.step(timestep) != -1):
    if robot.getTime() > 1.0:
        break;
def doNothing(x):
    pass
# credit to Khawar Jamil for the openCV trackbar useage and object detection: https://medium.com/globant/maneuvering-color-mask-into-object-detection-fce61bf891d1
cv2.namedWindow('Track Bars', cv2.WINDOW_NORMAL)
cv2.createTrackbar('min_blue', 'Track Bars', 0, 255, doNothing)
cv2.createTrackbar('min_green', 'Track Bars', 0, 255, doNothing)
cv2.createTrackbar('min_red', 'Track Bars', 0, 255, doNothing)

cv2.createTrackbar('max_blue', 'Track Bars', 0, 255, doNothing)
cv2.createTrackbar('max_green', 'Track Bars', 0, 255, doNothing)
cv2.createTrackbar('max_red', 'Track Bars', 0, 255, doNothing)
Motor.setPosition(camera_pitch_motor, 1.5);
obj = 'goal'
while (robot.step(timestep) != -1):
    data = camera.getImage()
    if data:
        image = np.frombuffer(data, np.uint8).reshape((camera.getHeight(), camera.getWidth(), 4))
        img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        if obj is 'test':
            min_blue = cv2.getTrackbarPos('min_blue', 'Track Bars')
            min_green = cv2.getTrackbarPos('min_green', 'Track Bars')
            min_red = cv2.getTrackbarPos('min_red', 'Track Bars')
            max_blue = cv2.getTrackbarPos('max_blue', 'Track Bars')
            max_green = cv2.getTrackbarPos('max_green', 'Track Bars')
            max_red = cv2.getTrackbarPos('max_red', 'Track Bars')
        if obj is 'goal':
            min_blue = 54
            min_green = 179
            min_red = 232
            max_blue = 71
            max_green = 255
            max_red = 236
        if obj is 'obstacles':
            min_blue = 116
            min_green = 186
            min_red = 226
            max_blue =144
            max_green =255
            max_red =255
        mask = cv2.inRange(hsv, (min_blue, min_green, min_red), (max_blue, max_green, max_red))
        cv2.imshow('Mask Image', mask)
        key = cv2.waitKey(timestep)
        if mask is not None:
            #credit to Opencv documentation for contour checking and position approximation: https://docs.opencv.org/3.1.0/df/d9d/tutorial_py_colorspaces.html
            contours,hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
            if contours:
                for cnt in contours:
                    M = cv2.moments(cnt)
                    if M['m00'] != 0:
                        cx = int(M['m10']/M['m00'])
                        cy = int(M['m01']/M['m00'])
                        print(cx,cy)
                    area = cv2.contourArea(cnt)
                    perimeter = cv2.arcLength(cnt,True)
                    epsilon = 0.1*cv2.arcLength(cnt,True)
                    approx = cv2.approxPolyDP(cnt,epsilon,True)
                    x,y,w,h = cv2.boundingRect(cnt)
                    cv2.rectangle(image,(x,y),(x+w,y+h),(0,255,0),2)
                cv2.imshow("preview", image)
                cv2.waitKey(timestep)
    time = robot.getTime();  # in seconds.

    roll = InertialUnit.getRollPitchYaw(imu)[0] + np.pi / 2.0;
    pitch = InertialUnit.getRollPitchYaw(imu)[1];
    altitude = GPS.getValues(gps)[1];
    roll_acceleration = Gyro.getValues(gyro)[0];
    pitch_acceleration = Gyro.getValues(gyro)[1];
    num_objects = camera.getRecognitionNumberOfObjects()
    objects = camera.getRecognitionObjects()
    # Stabilize the Camera by actuating the camera motors according to the gyro feedback.
    Motor.setPosition(camera_roll_motor, -0.115 * roll_acceleration);
    Motor.setPosition(camera_pitch_motor, 1.5);

    roll_disturbance = 0.0;
    pitch_disturbance = 0.0;
    yaw_disturbance = 0.0;
    key = keyboard.getKey();

    while (key > 0):
        if(key == Keyboard.UP):
            pitch_disturbance = 2.0;
            break;
        elif key==Keyboard.DOWN:
          pitch_disturbance = -2.0;
          break;
        elif key==Keyboard.RIGHT:
          yaw_disturbance = 1.3;
          break;
        elif key==Keyboard.LEFT:
          yaw_disturbance = -1.3;
          break;
        elif (key == Keyboard.SHIFT + Keyboard.RIGHT):
          roll_disturbance = -1.0;
          break;
        elif (key == Keyboard.SHIFT + Keyboard.LEFT):
          roll_disturbance = 1.0;
          break;
        elif (key == Keyboard.SHIFT + Keyboard.UP):
          target_altitude  += 0.05;
          print("target altitude: %f [m]\n", target_altitude);
          break;
        elif (key == Keyboard.SHIFT + Keyboard.DOWN):
          target_altitude -= 0.05;
          print("target altitude: %f [m]\n", target_altitude);
          break;
    key = keyboard.getKey();

    #Compute the roll, pitch, yaw and vertical inputs.
    roll_input = k_roll_p * CLAMP(roll, -1.0, 1.0) + roll_acceleration + roll_disturbance;
    pitch_input = k_pitch_p * CLAMP(pitch, -1.0, 1.0) - pitch_acceleration + pitch_disturbance;
    yaw_input = yaw_disturbance;
    clamped_difference_altitude = CLAMP(target_altitude - altitude, -1, 1);
    vertical_input = k_vertical_p * pow(clamped_difference_altitude, 3.0);

    # Actuate the motors taking into consideration all the computed inputs.
    front_left_motor_input = k_vertical_thrust + vertical_input - roll_input - pitch_input + yaw_input;
    front_right_motor_input = k_vertical_thrust + vertical_input + roll_input - pitch_input - yaw_input;
    rear_left_motor_input = k_vertical_thrust + vertical_input - roll_input + pitch_input - yaw_input;
    rear_right_motor_input = k_vertical_thrust + vertical_input + roll_input + pitch_input + yaw_input;
    Motor.setVelocity(front_left_motor, front_left_motor_input);
    Motor.setVelocity(front_right_motor, -front_right_motor_input);
    Motor.setVelocity(rear_left_motor, -rear_left_motor_input);
    Motor.setVelocity(rear_right_motor, rear_right_motor_input);

wb_robot_cleanup();


from PyQt5.QtCore import Qt, QRectF
from PyQt5.QtGui import QIntValidator, QColor
from PyQt5.QtWidgets import QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QGraphicsView, QGraphicsScene, QCheckBox, \
    QMessageBox

class ScreenshotAreaLayout:
    def __init__(self, config, main_window, parent_layout):
        self.config = config
        self.main_window = main_window
        self.parent_layout = parent_layout

    def add_layout(self):
        screenshot_area_layout = QVBoxLayout()
        screenshot_area_layout.setObjectName("screenshot_area_layout")

        resolution_layout = QHBoxLayout()
        self.screenshot_area_title_label = QLabel("识别范围设置")
        self.screenshot_area_title_label.setAlignment(Qt.AlignCenter)
        self.screenshot_area_label = QLabel("识别区域：", self.main_window)
        self.screenshot_area_x_label = QLabel("x", self.main_window)

        self.width_input = QLineEdit(self.main_window)
        self.height_input = QLineEdit(self.main_window)
        self.width_input.setText(str(int(self.config.shot_width)))
        self.height_input.setText(str(int(self.config.shot_height)))
        # 连接信号和槽
        self.width_input.textChanged.connect(self.update_inner_rect_size)
        self.height_input.textChanged.connect(self.update_inner_rect_size)
        self.width_input.setValidator(QIntValidator(0, self.config.desktop_width))
        self.height_input.setValidator(QIntValidator(0, self.config.desktop_height))
        resolution_layout.addWidget(self.screenshot_area_label)
        resolution_layout.addWidget(self.width_input)
        resolution_layout.addWidget(self.screenshot_area_x_label)
        resolution_layout.addWidget(self.height_input)

        aim_radius_layout = QHBoxLayout()
        self.mouse_moving_radius_label = QLabel("腰射自瞄半径：")
        self.mouse_moving_radius_input = QLineEdit(self.main_window)
        self.mouse_moving_radius_input.setObjectName("mouse_moving_radius")
        self.mouse_moving_radius_input.setText(str(int(self.config.mouse_moving_radius)))
        self.mouse_moving_radius_input.textChanged.connect(self.update_inner_circle_size)

        self.aim_mouse_moving_radius_label = QLabel("瞄准自瞄半径：")
        self.aim_mouse_moving_radius_input = QLineEdit(self.main_window)
        self.aim_mouse_moving_radius_input.setObjectName("aim_mouse_moving_radius")
        self.aim_mouse_moving_radius_input.setText(str(int(self.config.aim_mouse_moving_radius)))
        # 连接信号和槽
        self.aim_mouse_moving_radius_input.textChanged.connect(self.update_inner_circle_size)

        aim_radius_layout.addWidget(self.mouse_moving_radius_label)
        aim_radius_layout.addWidget(self.mouse_moving_radius_input)
        aim_radius_layout.addWidget(self.aim_mouse_moving_radius_label)
        aim_radius_layout.addWidget(self.aim_mouse_moving_radius_input)

        multi_stage_aiming_speed_layout = QHBoxLayout()
        self.multi_stage_aiming_speed_label = QLabel("腰射多级瞄速：")
        self.multi_stage_aiming_speed_input = QLineEdit(self.main_window)
        self.multi_stage_aiming_speed_input.setObjectName("multi_stage_aiming_speed")

        # 连接信号和槽
        multi_stage_aiming_speed_layout.addWidget(self.multi_stage_aiming_speed_label)
        multi_stage_aiming_speed_layout.addWidget(self.multi_stage_aiming_speed_input)

        aim_multi_stage_aiming_speed_layout = QHBoxLayout()
        self.aim_multi_stage_aiming_speed_label = QLabel("瞄准多级瞄速：")
        self.aim_multi_stage_aiming_speed_input = QLineEdit(self.main_window)
        self.aim_multi_stage_aiming_speed_input.setObjectName("aim_multi_stage_aiming_speed")

        # 连接信号和槽
        aim_multi_stage_aiming_speed_layout.addWidget(self.aim_multi_stage_aiming_speed_label)
        aim_multi_stage_aiming_speed_layout.addWidget(self.aim_multi_stage_aiming_speed_input)

        self.view = RectView(self.main_window,
                             outer_rect_size=(
                                 int(self.config.desktop_width / 10), int(self.config.desktop_height / 10)),
                             inner_rect_size=(
                                 int(self.config.shot_width / 10), int(self.config.shot_height / 10)),
                             radius=int(self.config.mouse_moving_radius / 10),
                             aim_radius=int(self.config.aim_mouse_moving_radius / 10))
        screenshot_area_layout.addWidget(self.screenshot_area_title_label)
        screenshot_area_layout.addLayout(resolution_layout)
        screenshot_area_layout.addLayout(aim_radius_layout)
        screenshot_area_layout.addLayout(multi_stage_aiming_speed_layout)
        screenshot_area_layout.addLayout(aim_multi_stage_aiming_speed_layout)
        screenshot_area_layout.addWidget(self.view)
        self.parent_layout.addLayout(screenshot_area_layout)

        self.init_form_config()

    def init_form_config(self):
        self.width_input.setText(str(int(self.config.shot_width)))
        self.height_input.setText(str(int(self.config.shot_height)))
        self.width_input.setValidator(QIntValidator(0, self.config.desktop_width))
        self.height_input.setValidator(QIntValidator(0, self.config.desktop_height))
        self.mouse_moving_radius_input.setText(str(int(self.config.mouse_moving_radius)))
        self.aim_mouse_moving_radius_input.setText(str(int(self.config.aim_mouse_moving_radius)))
        self.multi_stage_aiming_speed_input.setText(
            " ".join([f"{start}-{end}" for start, end in self.config.multi_stage_aiming_speed]))
        self.aim_multi_stage_aiming_speed_input.setText(
            " ".join([f"{start}-{end}" for start, end in self.config.aim_multi_stage_aiming_speed]))

    def update_inner_rect_size(self):
        # 当输入框的内容改变时，更新内部框的大小
        width = int(self.width_input.text()) if self.width_input.text() else 0
        height = int(self.height_input.text()) if self.height_input.text() else 0
        self.view.resize_inner_rect(width, height)

    def update_inner_circle_size(self):
        object_name = self.main_window.sender().objectName()
        if object_name == "mouse_moving_radius":
            radius = int(self.mouse_moving_radius_input.text()) if self.mouse_moving_radius_input.text() else 0
            self.view.resize_inner_circle(radius)
        elif object_name == "aim_mouse_moving_radius":
            radius = int(self.aim_mouse_moving_radius_input.text()) if self.aim_mouse_moving_radius_input.text() else 0
            self.view.resize_inner_circle_aim(radius)

    def check_multi_stage_aiming_speed(self, speed_up, multi_stage_aiming_speed_str):
        if multi_stage_aiming_speed_str is None or multi_stage_aiming_speed_str == "":
            return []
        multi_stage_aiming_speed_arr = multi_stage_aiming_speed_str.split(" ")
        number_array = []
        for num_str in multi_stage_aiming_speed_arr:
            try:
                num_str_arr = num_str.split("-")
                num_one = int(num_str_arr[0])
                num_two = int(num_str_arr[1])
                if not (len(num_str_arr) == 2 and num_two >= num_one):
                    QMessageBox.warning(self.main_window, "不符合条件",
                                        f"{num_str_arr} 格式错误，格式为 数字-数字，且前一位大于后一位")

                if not 0 <= num_two <= speed_up:
                    QMessageBox.warning(self.main_window, "不符合条件", f"{num_two} 数字不允许比瞄准范围大")
                    return []
                else:
                    number_array.append((num_one, num_two))
            except ValueError:
                QMessageBox.critical(self.main_window, "错误", f"{num_str} 格式错误")
                return []
        return number_array

    def save_config(self):
        self.config.set_config("shot_width", int(self.view.inner_rect.rect().width() * 10))
        self.config.set_config("shot_height", int(self.view.inner_rect.rect().height() * 10))
        self.config.set_config("mouse_moving_radius", int(self.mouse_moving_radius_input.text()))
        self.config.set_config("aim_mouse_moving_radius", int(self.aim_mouse_moving_radius_input.text()))

        multi_stage_aiming_speed_arr = self.check_multi_stage_aiming_speed(self.config.mouse_moving_radius,
                                                                           self.multi_stage_aiming_speed_input.text())
        self.config.set_config("multi_stage_aiming_speed", multi_stage_aiming_speed_arr)

        aim_multi_stage_aiming_speed_arr = self.check_multi_stage_aiming_speed(self.config.aim_mouse_moving_radius,
                                                                               self.aim_multi_stage_aiming_speed_input
                                                                               .text())
        self.config.set_config("aim_multi_stage_aiming_speed", aim_multi_stage_aiming_speed_arr)

class RectView(QGraphicsView):
    def __init__(self, parent=None, outer_rect_size=(192, 108), inner_rect_size=(64, 64), radius=32, aim_radius=32):
        super(RectView, self).__init__(parent)
        self.setMinimumSize(*outer_rect_size)
        self.setScene(QGraphicsScene(self))

        self.outer_rect = self.scene().addRect(QRectF(0, 0, *outer_rect_size))  # 外部框
        self.outer_rect.setBrush(QColor(255, 0, 0))

        self.inner_rect = self.scene().addRect(QRectF(0, 0, *inner_rect_size))  # 内部框
        self.inner_rect.setBrush(QColor(0, 255, 0))
        self.center_inner_rect()

        self.inner_circle = self.scene().addEllipse(QRectF(0, 0, radius * 2, radius * 2))
        self.inner_circle.setBrush(QColor(0, 0, 255))
        self.center_inner_circle()

        self.inner_circle_aim = self.scene().addEllipse(QRectF(0, 0, aim_radius * 2, aim_radius * 2))
        self.inner_circle_aim.setBrush(QColor(0, 255, 255))
        self.center_inner_circle_aim()

    def center_inner_rect(self):
        # 将内部框居中
        self.inner_rect.setPos((self.outer_rect.rect().width() - self.inner_rect.rect().width()) / 2,
                               (self.outer_rect.rect().height() - self.inner_rect.rect().height()) / 2)

    def center_inner_circle(self):
        self.inner_circle.setPos((self.outer_rect.rect().width() - self.inner_circle.rect().width()) / 2,
                                 (self.outer_rect.rect().height() - self.inner_circle.rect().height()) / 2)

    def center_inner_circle_aim(self):
        self.inner_circle_aim.setPos((self.outer_rect.rect().width() - self.inner_circle_aim.rect().width()) / 2,
                                     (self.outer_rect.rect().height() - self.inner_circle_aim.rect().height()) / 2)

    def resize_inner_rect(self, width, height):
        # 改变内部框的大小
        self.inner_rect.setRect(0, 0, width / 10, height / 10)
        self.center_inner_rect()

    def resize_inner_circle(self, radius):
        self.inner_circle.setRect(0, 0, radius * 2 / 10, radius * 2 / 10)
        self.center_inner_circle()

    def resize_inner_circle_aim(self, radius):
        self.inner_circle_aim.setRect(0, 0, radius * 2 / 10, radius * 2 / 10)
        self.center_inner_circle_aim()


import random

from absl.testing import parameterized
import numpy as np

from tensorflow.python.client import session
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import googletest
from tensorflow.python.util import nest

def make_placeholder(t):
  return array_ops.placeholder_with_default(t, None)

def rebuild_ragged_tensor_with_value_rowids(rt, feed_dict=None, sess=None):
  """Returns a copy of `rt`, built using `from_value_rowids`.

  This ensures that RaggedTensor._cached_value_rowids is populated, which
  triggers a different code-path for converting ragged tensors to tensors.

  If `feed_dict` and `sess` are specified, then build the new `RaggedTensor`
  using placeholder tensors, and populate a feed dictionary that can be used
  to feed the placeholders.

  Args:
    rt: The RaggedTensor to copy.
    feed_dict: If specified, then build the new `RaggedTensor` using
      placeholders, and populate this dict with entries to feed those
      placeholders.
    sess: A session used to evaluate tensors; required if feed_dict is
      specified.

  Returns:
    A copy of `rt`, built using `from_value_rowids`.
  """
  if isinstance(rt, ragged_tensor.RaggedTensor):
    values = rebuild_ragged_tensor_with_value_rowids(rt.values, feed_dict, sess)
    rowids = rt.value_rowids()
    nrows = rt.nrows()
    if feed_dict is not None:
      rowids_ph = make_placeholder(rowids)
      nrows_ph = make_placeholder(nrows)
      feed_dict[rowids_ph] = sess.run(rowids)
      feed_dict[nrows_ph] = sess.run(nrows)
      rowids, nrows = rowids_ph, nrows_ph
    return ragged_tensor.RaggedTensor.from_value_rowids(values, rowids, nrows)
  else:
    if feed_dict is not None:
      rt_ph = make_placeholder(rt)
      feed_dict[rt_ph] = sess.run(rt)
      rt = rt_ph
    return rt

@test_util.run_all_in_graph_and_eager_modes
class RaggedTensorToTensorOpTest(test_util.TensorFlowTestCase,
                                 parameterized.TestCase):

  def testDocStringExamples(self):
    """Example from ragged_to_tensor.__doc__."""
    rt = ragged_factory_ops.constant([[9, 8, 7], [], [6, 5], [4]])
    dt = rt.to_tensor()
    self.assertAllEqual(dt, [[9, 8, 7], [0, 0, 0], [6, 5, 0], [4, 0, 0]])

  @parameterized.named_parameters(
      # Simple 2D ragged tensors (with one ragged dimension)
      {
          'testcase_name': 'shape_2xN',
          'rt_input': [[0, 1, 2], [], [3]],
          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0]]
      },
      {
          'testcase_name': 'shape_2xN_default_0D',
          'rt_input': [[0, 1, 2], [], [3]],
          'default': 5,
          'expected': [[0, 1, 2], [5, 5, 5], [3, 5, 5]]
      },
      {
          'testcase_name': 'empty_first_row',
          'rt_input': [[], [], [3, 4], []],
          'expected': [[0, 0], [0, 0], [3, 4], [0, 0]]
      },
      {
          'testcase_name': 'empty_last_row',
          'rt_input': [[0, 1, 2], [], [3], []],
          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0], [0, 0, 0]]
      },
      {
          'testcase_name': 'shape_4xN',
          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
          'expected': [[1, 2, 3], [0, 0, 0], [4, 0, 0], [5, 6, 0]]
      },
      {
          'testcase_name': 'shape_4xN_default_0D',
          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
          'default': 9,
          'expected': [[1, 2, 3], [9, 9, 9], [4, 9, 9], [5, 6, 9]]
      },
      {
          'testcase_name': 'shape_2xN_already_dense',
          'rt_input': [[6, 7, 8], [9, 10, 11]],
          'expected': [[6, 7, 8], [9, 10, 11]],
      },
      {
          'testcase_name': 'shape_2xN_string_already_dense',
          'rt_input': [[b'a', b'b', b'c'],
                       [b'd', b'e', b'antidisestablishmentarianism']],
          'ragged_rank': 1,
          'expected': [[b'a', b'b', b'c'],
                       [b'd', b'e', b'antidisestablishmentarianism']],
      },
      # 3D ragged tensors with two ragged dimensions
      {
          'testcase_name': 'shape_4xNxM',
          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
          'expected': [
              [[1, 2], [0, 0], [3, 4]],  #
              [[0, 0], [0, 0], [0, 0]],  #
              [[5, 0], [0, 0], [0, 0]],  #
              [[6, 7], [8, 0], [0, 0]],  #
          ]
      },
      {
          'testcase_name': 'shape_4xNxM_default_0D',
          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
          'default': 9,
          'expected': [
              [[1, 2], [9, 9], [3, 4]],  #
              [[9, 9], [9, 9], [9, 9]],  #
              [[5, 9], [9, 9], [9, 9]],  #
              [[6, 7], [8, 9], [9, 9]],  #
          ]
      },
      {
          'testcase_name': 'shape_1xNx1_default_0D',
          'rt_input': [[[1], [2], [3]]],
          'ragged_rank': 1,
          'default': 0,
          'expected': [[[1], [2], [3]]],
      },
      {
          'testcase_name': 'shape_2xNx2_already_dense',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
          'ragged_rank': 1,
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
      },
      {
          'testcase_name': 'shape_2xNx2_already_dense_default_1D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
          'ragged_rank': 1,
          'default': [31, 32],
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
      },
      {
          'testcase_name': 'shape_2xNx2_string_already_dense',
          'rt_input': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
                        [b'manzana', b'nectar']]],
          'ragged_rank': 1,
          'expected': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
                        [b'manzana', b'nectar']]],
      },
      # 3D ragged tensors with one ragged dimension
      {
          'testcase_name': 'shape_4xNx1_default_1D',
          'rt_input': [[[1], [2], [3]], [], [[4]], [[5], [6]]],
          'ragged_rank': 1,
          'default': [9],
          'expected': [[[1], [2], [3]],
                       [[9], [9], [9]],
                       [[4], [9], [9]],
                       [[5], [6], [9]]]
      },
      {
          'testcase_name': 'shape_2xNx2_default_0D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15]]],
          'ragged_rank': 1,
          'default': 2,
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [2, 2]]],
      },
      {
          'testcase_name': 'shape_2xNx2_default_1D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15]]],
          'ragged_rank': 1,
          'default': [2, 3],
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [2, 3]]],
      },
      # 4D ragged tensors with 3 ragged dimensions
      {
          'testcase_name': 'shape_1xNxMxK_default_0D',
          'rt_input': [[[[1], [2]], [], [[3]]]],
          'default': 9,
          'expected': [[[[1], [2]], [[9], [9]], [[3], [9]]]],
      },
      # Broadcast default
      {
          'testcase_name': 'shape_2xNx2x2_default_2x1',
          'rt_input': [[[[1, 2], [3, 4]]], []],
          'ragged_rank': 1,
          'default': [[5], [6]],
          'expected': [[[[1, 2], [3, 4]]],
                       [[[5, 5], [6, 6]]]],
      },
      {
          'testcase_name': 'shape_2xNx2x2_default_1x2',
          'rt_input': [[[[1, 2], [3, 4]]], []],
          'ragged_rank': 1,
          'default': [[5, 6]],
          'expected': [[[[1, 2], [3, 4]]],
                       [[[5, 6], [5, 6]]]],
      },
      # Explicit shape
      {
          'testcase_name': 'shape_4xN_with_crop',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': [2, 3],
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_2xN_with_pad',
          'rt_input': [[1, 2], [3]],
          'shape': [3, 3],
          'expected': [[1, 2, 0], [3, 0, 0], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_crop_and_pad',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': [2, 8],
          'expected': [[0, 1, 2, 3, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_tuple_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': (2, 3),
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_tensorshape_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': tensor_shape.TensorShape([2, 3]),
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_partial_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': tensor_shape.TensorShape([2, None]),
          'expected': [[0, 1, 2, 3], [0, 0, 0, 0]],
      },
      # Empty tensors
      {
          'testcase_name': 'shape_0xN',
          'rt_input': [],
          'ragged_rank': 1,
          'expected': [],
          'expected_shape': [0, 0],
      },
      {
          'testcase_name': 'shape_0xNxM',
          'rt_input': [],
          'ragged_rank': 2,
          'expected': [],
          'expected_shape': [0, 0, 0],
      },
      # {
      #     'testcase_name': 'shape_0xNx2',
      #     'rt_input': [],
      #     'ragged_rank': 1,
      #     'inner_shape': [2],
      #     'expected': [],
      #     'expected_shape': [0, 0, 2],
      # },
      {
          'testcase_name': 'shape_2xN_empty',
          'rt_input': [[], []],
          'expected': [[], []],
          'expected_shape': [2, 0],
      },
  )  # pyformat: disable
  def testRaggedTensorToTensor(self,
                               rt_input,
                               expected,
                               ragged_rank=None,
                               inner_shape=None,
                               default=None,
                               shape=None,
                               expected_shape=None):
    rt1 = ragged_factory_ops.constant(
        rt_input, ragged_rank=ragged_rank, inner_shape=inner_shape)
    rt2 = rebuild_ragged_tensor_with_value_rowids(rt1)
    for rt in [rt1, rt2]:
      for use_placeholder in [False, True]:
        if use_placeholder:
          if default is not None:
            default = make_placeholder(default)
          rt = nest.map_structure(make_placeholder, rt, expand_composites=True)
        dt = rt.to_tensor(default_value=default, shape=shape)
        self.assertIsInstance(dt, tensor_lib.Tensor)
        self.assertEqual(rt.dtype, dt.dtype)
        if shape is not None:
          self.assertTrue(dt.shape.is_compatible_with(shape))
        else:
          self.assertTrue(dt.shape.is_compatible_with(rt.shape))
        if expected_shape is not None:
          expected = np.ndarray(expected_shape, buffer=np.array(expected))
        self.assertAllEqual(dt, expected)

  @parameterized.parameters([
      {
          'rt_input': [[1, 2, 3]],
          'default': 'a',
          'error_type': TypeError,
          'error': r'Expected int32|Cannot convert',
      },
      {
          'rt_input': [[1, 2, 3]],
          'default': [0],
          'error': r'default_value\.shape=.* and '
                   r'rt_input\.flat_values\.shape=.* are incompatible: '
                   r'default_value\.rank = 1  must be less than '
                   r'rt_input\.flat_values\.rank = 1'
      },
      {
          'rt_input': [[[1, 2], [3, 4]], [[5, 6]]],
          'ragged_rank': 1,
          'default': [7, 8, 9],
          'error': r'default_value\.shape.* and '
                   r'rt_input\.flat_values\.shape.* are incompatible: '
                   r'default_value\.shape\[-1\] = 3 but '
                   r'rt_input\.flat_values\.shape\[-1\] = 2'
      },
      {
          'rt_input': [[1, 2, 3]],
          'shape': [3, 3, 3],
          'error': r'rt_input\.shape and shape=\[.,.,.\] are incompatible: '
                   r'rt_input\.rank = 2 but shape\.rank = 3'
      },
      {
          'rt_input': [[[1, 2, 3]]],
          'ragged_rank': 1,
          'shape': [1, 1, 4],
          'error': r'rt_input\.shape and shape=\[1,1,4\] are incompatible: '
                   r'rt_input\.shape\[2\] = 3 but shape\[2\] = 4'
      },
  ])
  def testError(self,
                rt_input,
                error,
                error_type=(ValueError, errors.InvalidArgumentError),
                default=None,
                ragged_rank=None,
                shape=None):

    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
    with self.assertRaisesRegex(error_type, error):
      self.evaluate(rt.to_tensor(default_value=default, shape=shape))
    rt_placeholder = nest.map_structure(
        make_placeholder, rt, expand_composites=True)
    with self.assertRaisesRegex(error_type, error):
      self.evaluate(
          rt_placeholder.to_tensor(default_value=default, shape=shape))

  def test_shape_limit_shape_is_tensor(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, 3], dtype=dtypes.int64))
    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [2, 3])

  def test_shape_limit_shape_is_tensor_unknown_rank(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant(-1, dtype=dtypes.int64))
    self.assertAllEqual(
        actual, [[0, 1, 2, 3], [0, 0, 0, 0], [4, 0, 0, 0], [0, 0, 0, 0]])
    self.assertTrue(actual.shape.is_compatible_with([4, 4]))

  def test_shape_limit_shape_is_tensor_unknown_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, -1], dtype=dtypes.int64))
    self.assertAllEqual(actual, [[0, 1, 2, 3], [0, 0, 0, 0]])
    self.assertTrue(actual.shape.is_compatible_with([2, None]))

  def test_shape_limit_shape_is_tensor_int32(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, 3], dtype=dtypes.int32))
    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [2, 3])

  def test_shape_expand_first_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3]])
    actual = input_data.to_tensor(shape=[4, 4])
    self.assertAllEqual(
        actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0], [0, 0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [4, 4])

  def test_value_transposed(self):
    # Check that transposed data is not an issue.
    my_value = array_ops.transpose(
        constant_op.constant([[0, 1, 2, 3], [4, 5, 6, 7]]))
    input_data = RaggedTensor.from_value_rowids(
        values=my_value,
        value_rowids=constant_op.constant([0, 1, 2, 3], dtype=dtypes.int64),
        nrows=constant_op.constant(4, dtype=dtypes.int64),
        validate=True)
    self.assertAllEqual(input_data, [[[0, 4]], [[1, 5]], [[2, 6]], [[3, 7]]])

  def test_broadcast_default(self):
    # The dense dimension here is 2 x 2
    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
                                             ragged_rank=1)
    # This placeholder has a 2 x 1 dimension.
    default_value = make_placeholder([[5], [6]])
    actual = input_data.to_tensor(default_value=default_value)
    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
    self.assertAllEqual(actual, expected)

  def test_broadcast_default_no_placeholder(self):
    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
                                             ragged_rank=1)
    # default_value has a 2 x 1 dimension.
    default_value = constant_op.constant([[5], [6]], shape=None)
    actual = input_data.to_tensor(default_value=default_value)
    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
    self.assertAllEqual(actual, expected)

  def test_shape_expand_second_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3], []])
    actual = input_data.to_tensor(shape=[3, 4])
    self.assertAllEqual(actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0]])

  @parameterized.parameters(
      ([2, 3, 4], None, [2, 3, 4]),
      ([2, 3, 4], [None, None, None], [2, 3, 4]),
      ([2, 3, 4], [None, 3, None], [2, 3, 4]),
      ([2, 3, 4], [None, 3, 4], [2, 3, 4]),
      ([2, 3, 4], [2, 3, 4], [2, 3, 4]),
      )
  def test_preserve_shape_roundtrip(
      self, input_shape, to_tensor_shape, expected_shape):
    tensor = array_ops.zeros(input_shape)
    ragged_from_tensor = RaggedTensor.from_tensor(tensor, ragged_rank=2)
    recovered_tensor = ragged_from_tensor.to_tensor(shape=to_tensor_shape)
    self.assertAllEqual(tensor.shape.as_list(), expected_shape)
    self.assertAllEqual(ragged_from_tensor.shape.as_list(), expected_shape)
    self.assertAllEqual(recovered_tensor.shape.as_list(), expected_shape)

  def test_empty_tensor_with_shape(self):
    input_data = RaggedTensor.from_value_rowids(
        values=constant_op.constant([], dtype=dtypes.int64),
        value_rowids=constant_op.constant([], dtype=dtypes.int64),
        nrows=constant_op.constant(2, dtype=dtypes.int64),
        validate=True)
    actual = input_data.to_tensor(default_value=3, shape=[2, 3])
    self.assertAllEqual(actual, [[3, 3, 3], [3, 3, 3]])

  # pylint: disable=bad-whitespace
  @parameterized.named_parameters([
      dict(
          testcase_name = '2d_default_shape',
          shape         = None,
          rt_value      = [[1, 2, 3], [4], [5, 6]],
          rt_grad       = [[9, 8, 7], [6], [3, 2]],
          default_value = 0,
          default_grad  = sum([5, 4, 1]),
          output_value  = [[1, 2, 3], [4, 0, 0], [5, 6, 0]],
          output_grad   = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]),
      dict(
          testcase_name = '2d_pad',
          shape         = [4, 4],
          rt_value      = [[1, 2, 3], [4], [5, 6]],
          rt_grad       = [[9, 8, 7], [5], [1, 0]],
          default_value = 0,
          default_grad  = sum([6, 4, 3, 2, 1, 2, 3, 4, 5, 6]),
          output_value  = [
              [1, 2, 3, 0], [4, 0, 0, 0], [5, 6, 0, 0], [0, 0, 0, 0]],
          output_grad   = [
              [9, 8, 7, 6], [5, 4, 3, 2], [1, 0, 1, 2], [3, 4, 5, 6]]),
      dict(
          testcase_name = '2d_pad_and_crop',
          shape         = [5, 3],
          rt_value      = [[1, 2, 3], [4], [5, 6, 7, 8, 9], [8]],
          rt_grad       = [[9, 8, 7], [6], [3, 2, 1, 0, 0], [2]],
          default_value = 0,
          default_grad  = sum([5, 4, 3, 4, 5, 6, 7]),
          output_value  = [
              [1, 2, 3], [4, 0, 0], [5, 6, 7], [8, 0, 0], [0, 0, 0]],
          output_grad   = [
              [9, 8, 7], [6, 5, 4], [3, 2, 1], [2, 3, 4], [5, 6, 7]]),
      dict(
          testcase_name = '3d_rrank_2',
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8, 7], [6]], [[5, 4]]],
          rt_grad       = [[[1, 2, 0], [3]], [[5, 6]]],
          default_value = 3,
          default_grad  = sum([4, 7, 8]),
          output_value  = [[[9, 8], [6, 3]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_0d_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = 3,
          default_grad  = sum([7, 8]),
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_1d_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = [3, 2],
          default_grad  = [7, 8],
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 2]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_1d_broadcast_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = [3],
          default_grad  = [7 + 8],
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '4d_rrank_1_with_2d_default',
          ragged_rank   = 1,
          shape         = [3, 3, 2, 1],
          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
          default_value = [[3], [2]],
          default_grad  = [[5 + 9 + 2 + 4 + 6 + 8], [6 + 1 + 3 + 5 + 7 + 9]],
          output_value  = [[[[9], [8]], [[7], [6]], [[3], [2]]],
                           [[[5], [4]], [[3], [2]], [[3], [2]]],
                           [[[3], [2]], [[3], [2]], [[3], [2]]]],
          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
                           [[[7], [8]], [[9], [1]], [[2], [3]]],
                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
      dict(
          testcase_name = '4d_rrank_1_with_with_0d_default',
          ragged_rank   = 1,
          shape         = [3, 3, 2, 1],
          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
          default_value = 3,
          default_grad  = 5 + 9 + 2 + 4 + 6 + 8 + 6 + 1 + 3 + 5 + 7 + 9,
          output_value  = [[[[9], [8]], [[7], [6]], [[3], [3]]],
                           [[[5], [4]], [[3], [3]], [[3], [3]]],
                           [[[3], [3]], [[3], [3]], [[3], [3]]]],
          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
                           [[[7], [8]], [[9], [1]], [[2], [3]]],
                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
      dict(
          testcase_name = 'zero_size',
          shape         = [0, 0],
          rt_value      = [[9, 8], [7, 6, 5], [4]],
          rt_grad       = [[0, 0], [0, 0, 0], [0]],
          default_value = 3,
          default_grad  = 0,
          output_value  = [],
          output_grad   = [])
  ])  # pyformat: disable
  def test_gradient(self,
                    shape,
                    rt_value,
                    rt_grad,
                    default_value,
                    default_grad,
                    output_value,
                    output_grad,
                    ragged_rank=None):
    """Tests that ragged_to_dense generates the right gradient.

    Args:
      shape: The `shape` arg for `ragged_to_dense`.
      rt_value: The `rt_input` arg for `ragged_to_dense`.
      rt_grad: The expected gradient for `rt_value`.  Corresponds 1:1 with
        `rt_value`.
      default_value: The `default_value` arg for `ragged_to_dense`.
      default_grad: The expected gradient for `default_value`.  Corresponds 1:1
        with `default_value`.
      output_value: The expected output of `ragged_to_dense`.
      output_grad: The gradient for the output (used to generate the gradients
        `rt_grad` and `default_grad`).  Corresponds 1:1 with `output_value`.
      ragged_rank: Ragged rank for `rt_value`.
    """
    rt_value = ragged_factory_ops.constant(
        rt_value, dtype=dtypes.float32, ragged_rank=ragged_rank)
    rt_grad = ragged_factory_ops.constant(
        rt_grad, dtype=dtypes.float32, ragged_rank=ragged_rank)
    default_value = constant_op.constant(default_value, dtype=dtypes.float32)
    default_grad = constant_op.constant(default_grad, dtype=dtypes.float32)
    output_value = constant_op.constant(
        output_value, dtype=dtypes.float32, shape=shape)
    output_grad = constant_op.constant(
        output_grad, dtype=dtypes.float32, shape=shape)
    shape = tensor_shape.as_shape(shape)

    # There are different code paths for ragged_to_dense, depending on whether
    # the RaggedTensor was created from row_splits or value_rowids.  Make sure
    # that we test both.
    for partition_type in ['row_splits', 'value_rowids']:
      rt_val = self.rt_with_partition_type(rt_value, partition_type)
      if context.executing_eagerly():
        self._test_gradient_helper(rt_val, default_value, shape, output_grad,
                                   output_value, rt_grad, default_grad)
      else:
        # There are different code paths when computing the gradient for
        # default_value, depending on whether shape info is statically
        # available; make sure that we test all code paths.
        for shape_info in ['known', 'unknown_dims', 'unknown_rank']:
          rt_val = self.wrap_in_placeholder(rt_val, shape_info)
          default_val = self.wrap_in_placeholder(default_value, shape_info)
          shape_val = self.wrap_in_placeholder(shape, shape_info)
          self._test_gradient_helper(rt_val, default_val, shape_val,
                                     output_grad, output_value, rt_grad,
                                     default_grad)

  def _test_gradient_helper(self, rt_val, default_val, shape_val, output_grad,
                            expected_output_val, expected_rt_grad,
                            expected_default_grad):
    if context.executing_eagerly():
      with backprop.GradientTape() as tape:
        tape.watch([rt_val, default_val])
        out = rt_val.to_tensor(default_val, shape=shape_val)
        actual_rt_grad, actual_default_grad = tape.gradient(
            out, (rt_val, default_val), output_gradients=output_grad)
    else:
      out = rt_val.to_tensor(default_val, shape=shape_val)
      actual_rt_grad, actual_default_grad = gradients_impl.gradients(
          ys=out, xs=(rt_val, default_val), grad_ys=output_grad)

    self.assertAllClose(out, expected_output_val)
    self.assertIsInstance(actual_rt_grad, RaggedTensor)
    self.assertAllClose(actual_rt_grad, expected_rt_grad)
    self.assertAllClose(actual_default_grad, expected_default_grad)

  def rt_with_partition_type(self, rt, partition_type):
    if isinstance(rt, tensor_lib.Tensor):
      return rt
    if partition_type == 'row_splits':
      return rt
    if partition_type == 'value_rowids':
      return ragged_tensor.RaggedTensor.from_value_rowids(
          self.rt_with_partition_type(rt.values, partition_type),
          rt.value_rowids(), rt.nrows())
    raise AssertionError('Unexpected partition_type %r' % partition_type)

  def wrap_in_placeholder(self, arg, shape_info):
    """Wraps `arg` in a placeholder to limit static shape info.

    Args:
      arg: The value to wrap.  A Tensor, RaggedTensor, or TensorShape.
      shape_info: One of ['known', 'unknown_dims', 'unknown_rank'].

    Returns:
      * If shape_info is 'known': returns `arg`.
      * If shape_info is 'unknown_dims': returns a placeholder wrapping `arg`
        where the dimension sizes are unknown.  If `arg` is a TensorShape,
        then convert it to a vector first.  If `arg` is a RaggedTensor, then
        wrap the flat_values.
      * If shape_info is 'unknown_rank': returns a placeholder wrapping `arg`
        where the rank is unknown.  If `arg` is a TensorShape, then convert it
        to a vector first.  If `arg` is a RaggedTensor, then wrap the
        flat_values.
    """
    if shape_info == 'known':
      return arg
    if isinstance(arg, ragged_tensor.RaggedTensor):
      return arg.with_flat_values(
          self.wrap_in_placeholder(arg.flat_values, shape_info))
    if isinstance(arg, tensor_shape.TensorShape):
      if arg.ndims is None:
        return arg
      arg = constant_op.constant(arg.as_list())
    if shape_info == 'unknown_rank':
      return array_ops.placeholder_with_default(arg, None)
    if shape_info == 'unknown_dims':
      return array_ops.placeholder_with_default(arg, [None] * arg.shape.rank)
    raise AssertionError('Unexpected shape_info %r' % shape_info)

  def test_shape_is_list_including_tensor_element(self):
    rt = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6]])
    result = rt.to_tensor(shape=[2, constant_op.constant(2)])
    self.assertAllEqual(result, [[1, 2], [4, 0]])

class RaggedToDenseBenchmark(googletest.Benchmark):

  # Configurations to test.  See `run_benchmark` for config param docs.
  CONFIGS = [
      {'shape': [10, 10]},
      {'shape': [10, 1000]},
      {'shape': [1000, 10]},
      {'shape': [1000, 10], 'fill': [1, 0.95]},  # Mostly full.
      {'shape': [1000, 10], 'fill': [1, 0.05]},  # Mostly empty.
      {'shape': [1000, 10], 'dtype': dtypes.string},
      {'shape': [1000, 10], 'dtype': dtypes.int64},
      {'shape': [100, 100]},
      {'shape': [50, 50, 32]},
      {'shape': [100, 100, 100], 'min_iters': 100},
      {'shape': [1000, 1000], 'min_iters': 100},
      {'shape': [10, 10, 10, 10, 10]},
      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 1},
      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 2},
      {'shape': [50, 50, 32], 'ragged_rank': 1, 'default_shape': [32]},
      {'shape': [200, 50, 32], 'ragged_rank': 1, 'default_shape': [32]}
  ]  # pyformat: disable

  def run_benchmark(self,
                    shape=(100, 100),
                    ragged_rank=None,
                    dtype=dtypes.float32,
                    fill=None,
                    default_shape=(),
                    output_shape=None,
                    min_iters=1000):
    """Run a benchmark with the specified configuration parameters.

    Args:
      shape: Bounding box for the input ragged tensor.
      ragged_rank: Ragged rank for the input ragged tensor.  Defaults to
        `len(shape)-1`.
      dtype: Data type for the input ragged tensor.
      fill: How full each dimension should be (0-1).  Corresponds 1:1 with
        `shape`.  Defaults to 0.8 for each dimension.
      default_shape: Shape for the default (padding) value.
      output_shape: Output shape -- ragged tensor will be padded or cropped to
        this shape.
      min_iters: Minimum iterations for benchmark.
    """
    if ragged_rank is None:
      ragged_rank = len(shape) - 1
    if fill is None:
      fill = [0.8 for _ in shape]

    # Build the inputs for the op.
    rt_input = self._generateRaggedTensor(shape, ragged_rank, dtype, fill)
    default_value = constant_op.constant(
        self._generateRaggedTensor(default_shape, 0, dtype), dtype=dtype)

    mbs = np.prod(shape) / (2**20)
    with session.Session(config=benchmark.benchmark_config()) as sess:
      extras = {
          'shape': shape,
          'ragged_rank': ragged_rank,
          'dtype': dtype,
          'fill': fill,
          'default_shape': default_shape
      }
      rt = ragged_factory_ops.constant(rt_input, dtype, ragged_rank=ragged_rank)

      # Inputs for with_splits:
      splits_rt_placeholder = ragged_factory_ops.placeholder(
          dtype, ragged_rank, shape[ragged_rank + 1:])
      splits_feed_dict = {splits_rt_placeholder: sess.run(rt)}

      # Inputs for with_rowids:
      rowids_feed_dict = {}
      rowids_rt_placeholder = rebuild_ragged_tensor_with_value_rowids(
          rt, rowids_feed_dict, sess)

      # Common arguments for benchmarks:
      run_op_benchmark_kwargs = dict(
          sess=sess,
          store_memory_usage=True,
          min_iters=min_iters,
          burn_iters=max(5, min_iters // 10),
          mbs=mbs,
          extras=extras)

      ragged_to_tensor_with_splits = splits_rt_placeholder.to_tensor(
          default_value=default_value)
      self.run_op_benchmark(
          op_or_tensor=ragged_to_tensor_with_splits.op,
          name='ragged_to_tensor_with_splits',
          feed_dict=splits_feed_dict,
          **run_op_benchmark_kwargs)

      ragged_to_tensor_with_rowids = rowids_rt_placeholder.to_tensor(
          default_value=default_value)
      self.run_op_benchmark(
          op_or_tensor=ragged_to_tensor_with_rowids.op,
          name='ragged_to_tensor_with_rowids',
          feed_dict=rowids_feed_dict,
          **run_op_benchmark_kwargs)

  def _generateRaggedTensor(self, shape, ragged_rank, dtype, fill=None, axis=0):
    if axis == len(shape):
      value = random.random()
      if dtype == dtypes.string:
        value = str(value)
      if dtype.is_integer:
        value = int(value * 1000)
      return value
    if axis == 0 or axis > ragged_rank:
      slice_size = shape[axis]
    else:
      slice_size = (np.random.geometric(fill[axis], shape[axis]) == 1).sum()
    return [
        self._generateRaggedTensor(shape, ragged_rank, dtype, fill, axis + 1)
        for _ in range(slice_size)
    ]

  def benchmark_ragged_to_dense(self):
    random.seed(5)
    for config in self.CONFIGS:
      self.run_benchmark(**config)

if __name__ == '__main__':
  googletest.main()


import random
import time
import requests
import tweepy

def randomTimes(averagesecondsbetweencalls):
    print(random.gauss(averagesecondsbetweencalls, 60*1))
    return random.gauss(averagesecondsbetweencalls, 60*1)

#check all of followers
#if they have a star or a retweet, also retweet and/or star

access_token = "xxx"
access_token_secret = "xxx"
consumer_key = "xxx"
consumer_secret = "xxx"

def initTweepy():
    global the_api
    auth = tweepy.OAuthHandler(consumer_key, consumer_secret)
    auth.set_access_token(access_token, access_token_secret)
    the_api = tweepy.API(auth)

#creates a array
def getFriends():
    print("getFriends")
    global friendlist
    friendlist = the_api.followers_ids(id = "jakegosskuehn")
    return the_api.followers_ids(id = "jakegosskuehn")

#retweets a tweet
def retweetATweet(tweet_id):
    the_api.retweet(tweet_id)

#creates a favorite for a tweet
def createFavorite(tweet_id):
    ithe_api.create_favorite(tweet_id)

#use tweet.favorited boolean
#   tweet.retweeted boolean
#tweet.id is ID of the tweet, to avoid repeats
def getFavorites(user_id):
    initTweepy

#Gets all of the followers of the authoritating user
#Then gets their public tweets, and checks
    #if its favorited or RT'd, user will do the same.
    #then sleep for 3 min
def checkFriendTimelines():
    for friend_id in friendlist:
        isthetweetspecial = False
        public_tweets = the_api.user_timeline(id = friend_id )
        for tweet in public_tweets:
            try:
                if tweet.favorited is True:
                    isthetweetspecial = True
                    print("Favored")
                    createFavorite(tweet.id)

                if tweet.retweeted is True:
                    isthetweetspecial = True
                    print("RT")
                    tetweetATweet(tweet.id)
                if isthetweetspecial:
                    print(tweet.text)
            except UnicodeEncodeError as e:
                print ("encode error", e)

        print("Did one person, now sleeping")
        print(the_api.get_user(friend_id).screen_name)
        time.sleep(180)

def checkSpecificTweet():
    print(the_api.me)
    public_tweets = the_api.user_timeline()
    for tweet in public_tweets:
        print (tweet.id)
        print (tweet.text)
        print (tweet.favorited)
        time.sleep(10)

initTweepy()
getFriends()
checkSpecificTweet()

###checkFriendTimelines()


import base64
import json
import threading
import time
from io import BytesIO

import qrcode

from pyncm import load_response
from pyncm.apis.login import LoginQrcodeUnikey, \
    LoginQrcodeCheck, GetCurrentLoginStatus, write_login_info_and_dump_to_database, WriteLoginInfo, SetNewSession

def check_qr_status(uuid, username):
    """
    检查二维码扫描状态
    :param uuid: 二维码uuid
    :param username: 用户名
    :return:
    """
    break_time = 180
    while break_time:
        rsp = LoginQrcodeCheck(uuid)  # 检测扫描状态
        if rsp["code"] == 803:
            # 登录成功
            print(f"{rsp['code']} -- {rsp['message']}", "...")
            write_login_info_and_dump_to_database(GetCurrentLoginStatus(), username)
            # WriteLoginInfo(GetCurrentLoginStatus(), username)
            break
        elif rsp["code"] == 800:
            print('timeout')
            break
        time.sleep(1)
        break_time -= 1
    print('scan exited')

def get_qrcode(username):
    """
    获取登录二维码
    :param username: 用户名
    :return: 二维码base64文本
    """
    SetNewSession()
    uuid = LoginQrcodeUnikey()["unikey"]  # 获取 UUID
    print("UUID", uuid)
    url = f"https://music.163.com/login?codekey={uuid}"  # 二维码内容即这样的 URL
    img = qrcode.make(url)
    bytes_io = BytesIO()
    img.save(bytes_io, format='JPEG')
    th = threading.Thread(target=check_qr_status, args=[uuid, username])
    th.start()
    return str(base64.b64encode(bytes_io.getvalue()), 'utf-8')

def get_current_login_status(username):
    """
    从数据库获取并加载登录cookies返回用户信息
    :param username: 用户名
    :return: 用户信息dict
    """
    try:
        temp = load_response(username)
        response = temp['cloud']['response']
        cookie = temp['cloud']['cookie']
        WriteLoginInfo(response, cookie)
        return GetCurrentLoginStatus()
    except json.decoder.JSONDecodeError:
        return {}

if __name__ == '__main__':
    get_qrcode('')


# requires curl, and grib_copy for navgem (ECC grib_tools)

import re
import os
import subprocess
import time
import re
import requests

from datetime import datetime, timedelta
import sqlite3
import json
import shutil
# generalize to ecmwf (json), navgem(asc), and cmc
# todo need a space margin (1 gb at least, as 500 mb for one run)
# todo somehow need to be able to optionally download missing runs

import threading
import sys
import signal

# print based on level (lower is more important):
#  0 is None (for no messages printed)
#  1 is important (i.e. disk space, database error)
#  2 is notification level: messages for when download runs complete
#  3 is potentially not important (could not download files (as in an incomplete model run))
#  4 is for verbose debugging (print curl/grib_copy statements)
print_level_importance = 2

# how often to check (seconds)
backoff_time = 600
# wait between downloads (seconds)
sleep_time = 1
# number of last n runs to get (might get one more than this)
num_latest_runs = 4

# only download runs not already calculated
disturbances_db_file_path = '/home/db/Documents/JRPdata/cyclone-genesis/disturbances.db'

# GFS and NAV are in the same thread as they use the same server
model_names_to_download_thread_groupings = [['GFS', 'NAV'], ['CMC'], ['ECM']]

# use this to download individual grib files for NAVGEM so we can split them with grib_copy
tmp_download_dir = '/tmp'

# min space default is 100 MB (to make sure to not cause other programs to crash)
# tmp refers to tmp_download_dir, the rest refer to model_data_folders_by_model_name
min_space_bytes = {
    'tmp': '100000000',
    'GFS': '100000000',
    'ECM': '100000000',
    'CMC': '100000000',
    'NAV': '100000000',
}

model_data_folders_by_model_name = {
    'GFS': '/home/db/metview/JRPdata/globalmodeldata/gfs',
    'ECM': '/home/db/metview/JRPdata/globalmodeldata/ecm',
    'CMC': '/home/db/metview/JRPdata/globalmodeldata/cmc',
    'NAV': '/home/db/metview/JRPdata/globalmodeldata/nav'
}

url_folder_by_model = {
    'GFS': 'https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod',
    'ECM': 'https://data.ecmwf.int/forecasts',
    'CMC': 'http://hpfx.collab.science.gc.ca',
    'NAV': 'https://nomads.ncep.noaa.gov/pub/data/nccf/com/fnmoc/prod'
}

# the time steps in the file name are not all ### (ECM just uses a single digit)
model_time_step_str_format = {
    'GFS': '03',
    'ECM': '00',
    'CMC': '03',
    'NAV': '03'
}

# not including GFS analysis files (.anl)
total_model_time_steps = {
    'GFS': {
        '00': 65,
        '06': 65,
        '12': 65,
        '18': 65
    },
    'ECM': {
        '00': 41,
        '12': 41
    },
    'CMC': {
        '00': 41,
        '12': 41
    },
    'NAV': {
        '00': 31,
        '06': 31,
        '12': 31,
        '18': 31
    }
}
# only interested in 6 hour time steps
all_time_steps_by_model = {
    'GFS': {
        '00': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240, 246, 252, 258, 264, 270, 276, 282, 288, 294, 300, 306, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384],
        '06': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240, 246, 252, 258, 264, 270, 276, 282, 288, 294, 300, 306, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384],
        '12': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240, 246, 252, 258, 264, 270, 276, 282, 288, 294, 300, 306, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384],
        '18': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240, 246, 252, 258, 264, 270, 276, 282, 288, 294, 300, 306, 312, 318, 324, 330, 336, 342, 348, 354, 360, 366, 372, 378, 384]
        },
    'ECM': {
        '00': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240],
        '12': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240]
    },
    'CMC': {
        '00': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240],
        '12': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180, 186, 192, 198, 204, 210, 216, 222, 228, 234, 240]
    },
    'NAV': {
        '00': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180],
        '06': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180],
        '12': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180],
        '18': [0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96, 102, 108, 114, 120, 126, 132, 138, 144, 150, 156, 162, 168, 174, 180]
    }
}

expected_num_grib_files_by_model_name = {
    'GFS': 9,
    'CMC': 9,
    'ECM': 9,
    'NAV': 9,
    'UKM': 9
}

# in order to select grib files from a folder we need to ignore these file name extensions
ignored_file_extensions = [
    'asc',
    'idx',
    'json',
    'index'
]

# hours between model runs (not the time steps)
model_interval_hours = {
    'GFS': 6,
    'ECM': 12,
    'CMC': 12,
    'NAV': 6
}

# this is used to check for already downloaded files
time_step_re_str_by_model_name = {
    'GFS': r'.*?\.f(?P<time_step>\d\d\d)_',
    'ECM': r'.*?-(?P<time_step>\d+)h-oper-fc',
    'CMC': r'.*?\d+_P(?P<time_step>\d+)',
    'NAV': r'.*?navgem_\d{10}f(?P<time_step>\d\d\d)'
}

# Specifies what parameters to gather (and also used in how to name it for split gribs)
# GFS is what is used in the .idx file
# ECM is what is in the .index file (list of jsons)
# CMC is the portion of the file name
# NAV are the keys used for grib_copy (download the entire grib and split it)
params_by_model_name = {
    'GFS': [
        ':PRMSL:mean sea level:',
        ':HGT:250 mb:',
        ':HGT:850 mb:',
        ':UGRD:850 mb:',
        ':VGRD:850 mb:',
        ':UGRD:925 mb:',
        ':VGRD:925 mb:',
        ':UGRD:10 m above ground:',
        ':VGRD:10 m above ground:'
        ],
    'ECM': [
        {"levtype": 'sfc', "param": 'msl'},
        {"levtype": 'sfc', "param": '10u'},
        {"levtype": 'sfc', "param": '10v'},
        {"levtype": 'pl', "levlist": '250', "param": 'gh'},
        {"levtype": 'pl', "levlist": '850', "param": 'gh'},
        {"levtype": 'pl', "levlist": '850', "param": 'u'},
        {"levtype": 'pl', "levlist": '850', "param": 'v'},
        {"levtype": 'pl', "levlist": '925', "param": 'u'},
        {"levtype": 'pl', "levlist": '925', "param": 'v'},
    ],
    'CMC': [
        'PRMSL_MSL_0',
        'HGT_ISBL_250',
        'HGT_ISBL_850',
        'UGRD_ISBL_850',
        'UGRD_ISBL_925',
        'VGRD_ISBL_850',
        'VGRD_ISBL_925',
        'UGRD_TGL_10',
        'VGRD_TGL_10'
    ],
    'NAV': [
        {"levtype": 'sfc', "level": 0, "shortName": 'prmsl'},
        {"levtype": 'sfc', "level": 10, "shortName": '10u'},
        {"levtype": 'sfc', "level": 10, "shortName": '10v'},
        {"levtype": 'pl', "level": 250, "shortName": 'gh'},
        {"levtype": 'pl', "level": 850, "shortName": 'gh'},
        {"levtype": 'pl', "level": 850, "shortName": 'u'},
        {"levtype": 'pl', "level": 850, "shortName": 'v'},
        {"levtype": 'pl', "level": 925, "shortName": 'u'},
        {"levtype": 'pl', "level": 925, "shortName": 'v'}
    ]
}

# handle thread errors
exit_event = threading.Event()

# print lines at or below print_level_importance
def print_level(level, warning_str):
    if level <= print_level_importance:
        print(warning_str)

def have_enough_disk_space():
    paths_to_check = {}
    if tmp_download_dir:
        paths_to_check['tmp'] = tmp_download_dir

    for model_name, path in model_data_folders_by_model_name.items():
        paths_to_check[model_name] = path

    enough_space_for_all_paths = True
    for name, path in paths_to_check.items():
        try:
            dustats = shutil.disk_usage(path)
            if dustats:
                free_bytes = dustats.free
                if free_bytes < int(min_space_bytes[name]):
                    enough_space_for_all_paths = False
                    break

        except Exception as e:
            print_level(1, e)
            print_level(1, "Warning could not get disk space for {path}")

    return enough_space_for_all_paths

def get_calculated_and_missing_time_steps(model_name, model_timestamp):
    model_hour = f'{model_timestamp.hour:02}'
    # remove from this the calculated timesteps
    model_time_steps_missing = set(all_time_steps_by_model[model_name][model_hour])

    disturbance_candidates = get_disturbances_from_db(model_name, model_timestamp)
    time_steps_calculated = set()
    if disturbance_candidates:
        for time_step_str, candidates in disturbance_candidates.items():
            time_step_int = int(time_step_str)
            time_steps_calculated.add(time_step_int)
            # discard will not throw an error even if it is not there
            model_time_steps_missing.discard(time_step_int)
    return sorted(list(time_steps_calculated)), sorted(list(model_time_steps_missing))

def get_disturbances_from_db(model_name, model_timestamp):
    conn = None
    retrieved_data = {}
    try:
        # store disturbance candidates in database
        conn = sqlite3.connect(disturbances_db_file_path)
        cursor = conn.cursor()

        ds = model_timestamp.isoformat()

        cursor.execute('SELECT data FROM disturbances WHERE model_name = ? AND date = ?', (model_name, ds))
        result = cursor.fetchone()
        if result:
            retrieved_data = json.loads(result[0])

    except sqlite3.Error as e:
        print_level(1, f"SQLite error: {e}")
    finally:
        if conn:
            conn.close()

    return retrieved_data

# for interval_hours = 6, find 6 hour model timestamp prior to 00,06,12,18Z
def get_latest_possible_model_time(model_name):
    interval_hours = model_interval_hours[model_name]
    utc_now = datetime.utcnow()
    model_hour = utc_now.hour
    latest_hour = (model_hour // interval_hours) * interval_hours
    latest_time = datetime.replace(utc_now, hour = latest_hour, minute = 0, second = 0, microsecond = 0)
    return latest_time

# for interval_hours = 6, find 6 hour model timestamp prior to 00,06,12,18Z
def get_latest_n_model_times(model_name, n=num_latest_runs):
    model_times = []
    interval_hours = model_interval_hours[model_name]
    utc_now = datetime.utcnow()
    model_hour = utc_now.hour
    latest_hour = (model_hour // interval_hours) * interval_hours
    latest_time = datetime.replace(utc_now, hour = latest_hour, minute = 0, second = 0, microsecond = 0)

    model_times.append(latest_time)
    for n in range(n):
        latest_time = get_prior_model_time(model_name, latest_time)
        model_times.append(latest_time)
    return model_times

def get_prior_model_time(model_name, dt):
    interval_hours = model_interval_hours[model_name]
    prior_time = dt - timedelta(hours = interval_hours)
    return prior_time

def load_index_file_ecm(idx_file_path):
    with open(idx_file_path, 'r') as file:
        lines = file.readlines()
    return [json.loads(line) for line in lines]

def find_matching_lines_gfs(idx_file_path, params):
    matching_lines = []
    all_lines = []
    with open(idx_file_path, 'r') as file:
        lines = file.readlines()
        for i, line in enumerate(lines):
            if any(param in line for param in params):
                matching_lines.append((i, line))
            if ':' in line:
                all_lines.append((i, line))
    return all_lines, matching_lines

def generate_curl_grib_copy_commands_nav(timestamp_prefix, url_folder, url_base_file_name, output_dir, params):
    output_dir_timestamp = os.path.join(output_dir, timestamp_prefix)

    # have no index file so use requests.head first
    # this allows us to to check if the directory exists before trying to get all the split grib files
    try:
        res = requests.head(url_folder)
        if not res:
            print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
            return -1

        if res.status_code != 200:
            print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
            return -1
    except:
        print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
        return -1

    # download to tmp folder first
    output_file = os.path.join(tmp_download_dir, f"{timestamp_prefix}_{url_base_file_name}.grib2")

    # always overwrite temp file in case of corruption, since all the split gribs should be able to be produced if it is OK
    url = f'{url_folder}{url_base_file_name}.grib2'
    curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s {url} -o {output_file}"
    print_level(4, curl_command)

    # Run the subprocess and capture its output
    result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    time.sleep(sleep_time)
    if result.returncode != 0:
        print_level(3, f"Warning: Command failed with exit code {result.returncode}")
        print_level(3, f"Command: {curl_command}")
        print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
        return -2

    grib_file_path = output_file

    try:
        os.makedirs(output_dir_timestamp, exist_ok=True)
    except:
        print(1, f"Could not make folder: {output_dir_timestamp}")
        return -1

    for param in params:
        levtype = param["levtype"]
        level = param["level"]
        shortName = param["shortName"]

        output_file = os.path.join(output_dir_timestamp, f"{timestamp_prefix}_{url_base_file_name}_{shortName}_{levtype}_{level}.grib2")

        grib_copy_command = f"grib_copy -w levtype='{levtype}',level={level},shortName='{shortName}' '{grib_file_path}' '{output_file}'"
        print_level(4, grib_copy_command)

        # Run the subprocess and capture its output
        result = subprocess.run(grib_copy_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode != 0:
            print_level(3, f"Warning: Command failed with exit code {result.returncode}")
            print_level(3, f"Command: {curl_command}")
            print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
            # delete the tmp file now on an error
            try:
                os.remove(grib_file_path)
            except:
                # this is somewhat serious as we could run out of tmp space
                print_level(1, f"Warning could not delete tmp file: {grib_file_path}")

            return -2

    # delete the tmp file now that done
    try:
        os.remove(grib_file_path)
    except:
        print_level(1, f"Warning could not delete tmp file: {grib_file_path}")

    return 0

def generate_curl_commands_cmc(timestamp_prefix, url_folder, url_base_file_name, output_dir, params):
    output_dir_timestamp = os.path.join(output_dir, timestamp_prefix)

    # have no index file so use requests.head first
    # this allows us to to check if the directory exists before trying to get all the split grib files
    try:
        res = requests.head(url_folder)
        if not res:
            print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
            return -1

        if res.status_code != 200:
            print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
            return -1
    except:
        print_level(3, f"Warning: Could not get requests.head() from {url_folder}")
        return -1

    for param in params:
        url_base_file_name_with_param = url_base_file_name.replace(r'{PARAM}', param)

        output_file = os.path.join(output_dir_timestamp, f"{timestamp_prefix}_{url_base_file_name_with_param}.grib2")

        overwrite_or_download = True

        if os.path.exists(output_file):
            overwrite_or_download = False

        if overwrite_or_download:
            url = f'{url_folder}{url_base_file_name_with_param}.grib2'
            curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s {url} -o {output_file}"
            print_level(4, curl_command)

            # Run the subprocess and capture its output
            result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            time.sleep(sleep_time)
            if result.returncode != 0:
                print_level(3, f"Warning: Command failed with exit code {result.returncode}")
                print_level(3, f"Command: {curl_command}")
                print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
                return -2

    return 0

def generate_curl_commands_ecm(timestamp_prefix, idx_file_name, url_folder, url_base_file_name, output_dir, params):
    output_dir_timestamp = os.path.join(output_dir, timestamp_prefix)
    idx_file_path = os.path.join(output_dir, timestamp_prefix, f'{timestamp_prefix}_{idx_file_name}')
    #os.makedirs(output_dir_timestamp, exist_ok=True)
    if not os.path.exists(idx_file_path):
        url_idx = f'{url_folder}{idx_file_name}'
        curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s {url_idx} -o {idx_file_path}"
        print_level(4, curl_command)

        # Run the subprocess and capture its output
        result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        time.sleep(sleep_time)
        if result.returncode != 0:
            print_level(3, f"Warning: Command failed with exit code {result.returncode}")
            print_level(3, f"Command: {curl_command}")
            print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
            try:
                # clear directory if no files ready yet
                os.rmdir(output_dir_timestamp)
            except:
                pass

            return -1

    data = load_index_file_ecm(idx_file_path)

    for param_dict in params:
        levtype = param_dict["levtype"]
        levlist = param_dict.get("levlist", "")
        param = param_dict["param"]

        matching_lines = [entry for entry in data if entry["levtype"] == levtype and entry["param"] == param and entry.get("levelist", "") == levlist]

        for entry in matching_lines:
            start = int(entry["_offset"])
            length = int(entry["_length"])
            size = length
            end = start + length - 1

            param_name = entry["param"].replace(':', '_').replace(' ', '_').lower()
            if levlist == "":
                full_name = f'{param_name}_{levtype}'
            else:
                full_name = f'{param_name}_{levlist}_{levtype}'

            full_name = full_name.replace(':', '_').replace(' ', '_').lower()
            range_str = f"{start}-{end}"

            output_file = os.path.join(output_dir_timestamp, f"{timestamp_prefix}_{url_base_file_name}_{full_name}.grib2")

            overwrite_or_download = True

            if os.path.exists(output_file):
                if os.path.getsize(output_file) == size:
                    overwrite_or_download = False

            if overwrite_or_download:
                url = f'{url_folder}{url_base_file_name}.grib2'
                curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s -r {range_str} {url} -o {output_file}"
                print_level(4, curl_command)

                # Run the subprocess and capture its output
                result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                time.sleep(sleep_time)
                if result.returncode != 0:
                    print_level(3, f"Warning: Command failed with exit code {result.returncode}")
                    print_level(3, f"Command: {curl_command}")
                    print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
                    return -2

    return 0

def generate_curl_commands_gfs(timestamp_prefix, idx_file_name, url_folder, url_base_file_name, output_dir, params):
    output_dir_timestamp = os.path.join(output_dir, timestamp_prefix)
    idx_file_path = os.path.join(output_dir, timestamp_prefix, f'{timestamp_prefix}_{idx_file_name}')
    #os.makedirs(output_dir_timestamp, exist_ok=True)
    if not os.path.exists(idx_file_path):
        url_idx = f'{url_folder}{idx_file_name}'
        curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s {url_idx} -o {idx_file_path}"
        print_level(4, curl_command)

        # Run the subprocess and capture its output
        result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        time.sleep(sleep_time)
        if result.returncode != 0:
            print_level(3, f"Warning: Command failed with exit code {result.returncode}")
            print_level(3, f"Command: {curl_command}")
            print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
            try:
                # clear directory if no files ready yet
                os.rmdir(output_dir_timestamp)
            except:
                pass

            return -1

    all_lines, matching_lines = find_matching_lines_gfs(idx_file_path, params)
    total_lines = len(all_lines)

    for i, (line_num, line) in enumerate(matching_lines):
        fields = line.split(':')
        start = int(fields[1])
        if i == total_lines - 1:
            range_str = f"{start}-"
            # TODO, right now always overwrite last file? no, change from getting the http size of the url first and setting that as end-1
            size = -1
        else:
            next_line_fields = all_lines[line_num + 1][1].split(':')
            end = int(next_line_fields[1]) - 1
            range_str = f"{start}-{end}"
            size = end - start + 1

        param_name = fields[3].replace(':', '_').replace(' ', '_').lower()
        param_level = fields[4].replace(':', '_').replace(' ', '_').lower()
        output_file = os.path.join(output_dir_timestamp, f"{timestamp_prefix}_{os.path.basename(url_base_file_name)}_{param_name}_{param_level}.grib2")

        overwrite_or_download = True

        if os.path.exists(output_file):
            if os.path.getsize(output_file) == size:
                overwrite_or_download = False

        if overwrite_or_download:
            url = f'{url_folder}{url_base_file_name}'
            curl_command = f"curl --create-dirs -y 30 -Y 30 -f -v -s -r {range_str} {url} -o {output_file}"
            print_level(4, curl_command)

            # Run the subprocess and capture its output
            result = subprocess.run(curl_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            time.sleep(sleep_time)
            if result.returncode != 0:
                print_level(3, f"Warning: Command failed with exit code {result.returncode}")
                print_level(3, f"Command: {curl_command}")
                print_level(3, f"Error output: {result.stderr.decode('utf-8')}")
                return -2

    return 0

# the full timestep string as used in the file names we are parsing from the model data
# include any leading zeros up that make sense (only up to what the model covers)
def convert_model_time_step_to_str(model_name, model_time_step_int):
    str_format = model_time_step_str_format[model_name]
    return f'{model_time_step_int:{str_format}}'

def get_model_time_steps(model_name, model_timestamp):
    hour = model_timestamp.hour
    hour_str = f'{hour:02}'
    return all_time_steps_by_model[model_name][hour_str]

def get_completed_and_missing_downloaded_model_steps(model_name, model_timestamp):
    model_time_steps = get_model_time_steps(model_name, model_timestamp)
    model_base_dir = model_data_folders_by_model_name[model_name]
    model_date_str_with_hour = datetime.strftime(model_timestamp, '%Y%m%d%H')
    model_date_str = datetime.strftime(model_timestamp, '%Y%m%d')
    model_dir = os.path.join(model_base_dir, model_date_str_with_hour)
    model_hour = datetime.strftime(model_timestamp, '%H')
    model_time_step_re = re.compile(time_step_re_str_by_model_name[model_name])
    if not os.path.exists(model_dir):
        model_time_steps_int = [int(x) for x in model_time_steps]
        return [], model_time_steps_int
    files = os.listdir(model_dir)
    steps_downloaded = []
    steps_missing = []

    expected_num_grib_files_per_step = expected_num_grib_files_by_model_name[model_name]
    for model_time_step in model_time_steps:
        if type(model_time_step) is str:
            model_time_step_str = model
        elif type(model_time_step) is int:
            model_time_step_str = convert_model_time_step_to_str(model_name, model_time_step)

        # may not be a trailing slash or not
        if not os.path.isdir(model_dir):
            model_timestamp_str = os.path.basename(os.path.dirname(model_dir))
        else:
            model_timestamp_str = os.path.basename(model_dir)

        model_step_timestamp = datetime.strptime(model_timestamp_str, '%Y%m%d%H')

        grib_files = []

        for f in files:
            f_ext = f.split('.')[-1]
            if f_ext not in ignored_file_extensions:
                res = re.match(model_time_step_re, f)
                if res:
                    file_model_time_step_str = res['time_step']
                    if file_model_time_step_str == model_time_step_str:
                        grib_files.append(os.path.join(model_dir, f))

        num_grib_files_in_step = len(grib_files)
        if num_grib_files_in_step >= expected_num_grib_files_per_step:
            steps_downloaded.append(int(model_time_step_str))
        else:
            steps_missing.append(int(model_time_step_str))
    return steps_downloaded, steps_missing

# returns complete (boolean), downloaded steps (int list), missing steps (int list)
def model_run_steps_downloaded_info(model_name, model_timestamp):
    model_hour = datetime.strftime(model_timestamp, '%H')
    expected_num_total_steps = total_model_time_steps[model_name][model_hour]
    steps_downloaded, steps_missing = get_completed_and_missing_downloaded_model_steps(model_name, model_timestamp)

    return steps_downloaded, steps_missing

def download_step(model_name, model_timestamp, time_step_int):
    model_date = datetime.strftime(model_timestamp, '%Y%m%d')
    model_hour = datetime.strftime(model_timestamp, '%H')
    time_step = convert_model_time_step_to_str(model_name, time_step_int)

    output_dir = model_data_folders_by_model_name[model_name]
    # include trailing slash in url_folder (not in url_folder_by_model)
    model_url_folder = url_folder_by_model[model_name]
    # todo generalize:

    timestamp_prefix = f'{model_date}{model_hour}'
    params = params_by_model_name[model_name]

    if model_name == 'GFS':
        url_folder = f'{model_url_folder}/gfs.{model_date}/{model_hour}/atmos/'
        idx_file_name = f'gfs.t{model_hour}z.pgrb2.0p25.f{time_step}.idx'
        url_base_file_name = f'gfs.t{model_hour}z.pgrb2.0p25.f{time_step}'
        r = generate_curl_commands_gfs(timestamp_prefix, idx_file_name, url_folder, url_base_file_name, output_dir, params)
    elif model_name == 'ECM':
        url_folder = f'{model_url_folder}/{model_date}/{model_hour}z/0p4-beta/oper/'
        idx_file_name = f'{model_date}{model_hour}0000-{time_step}h-oper-fc.index'
        # we will add the .grib2 extension before we download
        url_base_file_name = f'{model_date}{model_hour}0000-{time_step}h-oper-fc'
        r = generate_curl_commands_ecm(timestamp_prefix, idx_file_name, url_folder, url_base_file_name, output_dir, params)
    elif model_name == 'CMC':
        url_folder = f'{model_url_folder}/{model_date}/WXO-DD/model_gem_global/15km/grib2/lat_lon/{model_hour}/{time_step}/'
        # we will add the .grib2 extension before we download and substitute PARAM ourselves
        url_base_file_name = r'CMC_glb_{PARAM}' + f'_latlon.15x.15_{model_date}{model_hour}_P{time_step}'
        r = generate_curl_commands_cmc(timestamp_prefix, url_folder, url_base_file_name, output_dir, params)
    elif model_name == 'NAV':
        url_folder = f'{model_url_folder}/navgem.{model_date}/'
        # we will add the .grib2 extension before we download and substitute PARAM ourselves
        url_base_file_name = f'navgem_{model_date}{model_hour}f{time_step}'
        r = generate_curl_grib_copy_commands_nav(timestamp_prefix, url_folder, url_base_file_name, output_dir, params)

    return r

def download_model_run(model_name, model_timestamp):
    already_calculated_time_steps, missing_calculated_time_steps = get_calculated_and_missing_time_steps(model_name, model_timestamp)

    if missing_calculated_time_steps is None:
        return 0

    steps_downloaded, steps_missing = model_run_steps_downloaded_info(model_name, model_timestamp)

    # remove already downloaded, calculated, and deleted steps
    steps_missing = sorted(list(set(steps_missing).intersection(set(missing_calculated_time_steps))))

    none_downloaded = (len(steps_downloaded) == 0)
    none_could_download = False
    r = 0
    if steps_missing:
        for time_step_int in steps_missing:
            if not have_enough_disk_space():
                print_level(1, "Fatal error: Exiting since disk space is below the minimum set.")
                exit()
            cur_r = download_step(model_name, model_timestamp, time_step_int)
            if cur_r != 0:
                if r != -1:
                    # OK or could not get one of the gribs
                    r = cur_r
                elif cur_r == -1:
                    # could not get an index
                    r = cur_r

                if none_downloaded:
                    none_could_download = True
                    r = -3
                    break
                # assume all time steps are complete (if there are missing time steps this will fail to download rest)
                break
            else:
                none_could_download = False
                none_downloaded = False
    return r

def download_thread(model_names):
    try:
        model_timestamps = {}
        complete_model_runs = {}
        for model_name in model_names:
            complete_model_runs[model_name] = []

        while not exit_event.is_set():
            exiting = False
            for model_name in model_names:
                model_timestamps[model_name] = get_latest_n_model_times(model_name)
                for model_timestamp in model_timestamps[model_name]:
                    if exit_event.is_set():
                        exiting = True
                        break
                    if model_timestamp in complete_model_runs[model_name]:
                        print_level(3, f'Already have {model_timestamp} for {model_name}.')
                        r = 0
                        continue
                    # download latest run with available timesteps
                    # there may be some cycling but the staleness check should always avoid infinite cycling between old runs
                    r = download_model_run(model_name, model_timestamp)
                    if r == -2 or r == -1:
                        print_level(3, f'Failed to finish downloading latest {model_name} run from {model_timestamp}. Retrying later.')
                    elif r == -3:
                        # could not even download any files, try the prior run
                        print_level(3, f'{model_name} run from {model_timestamp} not available.')
                    elif r == 0:
                        # successfully downloaded the latest run
                        print_level(2, f'Finished downloading latest {model_name} run for {model_name} {model_timestamp} (or already exists).')
                        # get the next one after backing off
                        if model_timestamp not in complete_model_runs[model_name]:
                            complete_model_runs[model_name].append(model_timestamp)
                if exiting:
                    break
            if exiting:
                break

            print_level(3, f'Backing off for {backoff_time} seconds in thread {threading.current_thread().name}.')
            time.sleep(backoff_time)

        if exiting:
            print_level(1, f'Thread {threading.current_thread().name} exiting.')

    except Exception as e:
        print_level(1, f"Thread {threading.current_thread().name} encountered an error: {e}")
        exit_event.set()
        sys.exit(1)

# Register a signal handler to catch Ctrl+C and exit gracefully
def signal_handler(sig, frame):
    print_level(1, f"\nCtrl+C received. Exiting once download runs are finished. May take {backoff_time} s to exit. Repeat to forcefully exit (may corrupt downloads)")
    exit_event.set()
    global ctrl_c_count
    ctrl_c_count += 1

    if ctrl_c_count >= 4:
        print("Received multiple Ctrl-C. Exiting forcefully immediately.")
        os._exit(1)

if __name__ == '__main__':
    # run the downloads per thread (default )
    threads = []

    ctrl_c_count = 0

    signal.signal(signal.SIGINT, signal_handler)

    for group in model_names_to_download_thread_groupings:
        thread = threading.Thread(target=download_thread, args=(group,), name='_'.join(group))
        thread.start()
        threads.append(thread)

    for thread in threads:
        thread.join()

    print_level(1, "All threads have finished.")


from typing import Generator, NamedTuple, Optional
from jaxtyping import Array, Float, Int

import jax.numpy as jnp
import jax.random as jr
import numpy as np

class Batch(NamedTuple):
    x: jnp.ndarray
    y: jnp.ndarray

def num_batches(x: jnp.ndarray, batch_size: int) -> int:
    """
    Compute the number of batches in an array.

    Args:
        x: The array to compute the number of batches in.
        batch_size: The size of each batch.

    Returns:
        The number of batches in the array.
    """
    num_complete_batches, leftover = divmod(x.shape[0], batch_size)
    return num_complete_batches + bool(leftover)

def num_batches(arr: Array, batch_size: int, axis: int = 0) -> int:
    num_full, leftover = divmod(arr.shape[axis], batch_size)
    return num_full + bool(leftover)

def data_loader(
    x: Float[Array, "n x1 x2 ch"],
    y: Int[Array, "n"],
    batch_size: int,
    shuffle: Optional[jr.KeyArray] = None,
):
    perm = np.arange(x.shape[0])
    if shuffle is not None:
        np.random.shuffle(perm)

    for i in range(num_batches(x, batch_size)):
        batch_idx = perm[i * batch_size : (i + 1) * batch_size]
        yield Batch(x[batch_idx], y[batch_idx])

class DataLoader:
    def __init__(
        self,
        x: jnp.ndarray,
        y: jnp.ndarray,
        batch_size: int,
        shuffle: Optional[jr.KeyArray] = None,
    ):
        self.x = x
        self.y = y
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self) -> Generator[Batch, None, None]:
        return data_loader(self.x, self.y, self.batch_size, self.shuffle)

    def __len__(self) -> int:
        return num_batches(self.x, self.batch_size)

    def __getitem__(self, idx: int) -> Batch:
        if idx >= len(self):
            raise IndexError
        return Batch(self.x[idx], self.y[idx])