import os
import torch

import numpy as np
from torch.distributions import Normal
import matplotlib.pyplot as plt



seed = 2023
device = f"cuda:0" if torch.cuda.is_available() else "cpu"


def generate_data(num, device="cpu"):
  
    theta = np.linspace(0, 2 * np.pi, num)
    radius = 1.0
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)
    # print(x)

   
    np.random.seed(2023)  
    noise = np.random.uniform(0, 0.05, x.shape)
    x_noisy = x + noise
    y_noisy = y + noise


    plt.figure(figsize=(6, 6))
    plt.scatter(x_noisy, y_noisy, c="royalblue", marker=".", label="noised_Circle")
    # sc = plt.scatter(x_noisy, y_noisy , c=theta, cmap=cm.viridis, marker='.', label='Circle')

    action_x = x_noisy.reshape(-1, 1)
    action_y = y_noisy.reshape(-1, 1)
    action = np.concatenate((action_x, action_y), axis=1)
    reward = np.ones_like(action_x)
    # print(reward.reshape(-1,1))
    state = np.zeros_like(action)
    # print(state)
    toy_dict = {
        "states": torch.from_numpy(state).float(),
        "next_states": torch.from_numpy(state).float(),
        "actions": torch.from_numpy(action).float(),
        "rewards": torch.from_numpy(reward.reshape(-1, 1)).float(),
        "done": torch.from_numpy(np.ones_like(reward).reshape(-1, 1)).float(),
        "is_finished": torch.from_numpy(np.ones_like(reward).reshape(-1, 1)).float(),
    }
    np.save("toy.npy", toy_dict)

    # return Data_Sampler( torch.from_numpy(state).float(),  torch.from_numpy(action).float(),  torch.from_numpy(reward.reshape(-1,1)).float(), device)


# seed = 2023
np.random.seed(seed)

num_data = int(1000)
generate_data(num_data, device)
