#Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.

import matplotlib.pyplot as plt
import os

import matplotlib.image as mpimg
import random
import math
import numpy as np
import pandas as pd 
if not os.path.exists('/home/dyr/fd/DT-GNF_pendulum/causal_data/pendulum/'): 
  os.makedirs('/home/dyr/fd/DT-GNF_pendulum/causal_data/pendulum/train/')
  os.makedirs('/home/dyr/fd/DT-GNF_pendulum/causal_data/pendulum/test/')

def projection(theta, phi, x, y, base = -0.5):
    b = y-x*math.tan(phi)
    shade = (base - b)/math.tan(phi)
    return shade
# 
scale = np.array([[0,44],[100,40],[7,7.5],[10,10]])#盲猜是坐标
count = 0
empty = pd.DataFrame(columns=['i', 'j', 'shade','mid'])
for i in np.linspace(-44,44,177):#pendulum（灯吧。）
    for j in np.linspace(60,140,161):#light（钟摆吧）
        if j == 100:
            continue
        plt.rcParams['figure.figsize'] = (1.0, 1.0)
        theta = i*math.pi/200.0
        phi = j*math.pi/200.0
        x = 10 + 8*math.sin(theta)
        y = 10.5 - 8*math.cos(theta)

        ball = plt.Circle((x,y), 1.5, color = 'firebrick')
        gun = plt.Polygon(([10,10.5],[x,y]), color = 'black', linewidth = 3)

        light = projection(theta, phi, 10, 10.5, 20.5)
        sun = plt.Circle((light,20.5), 3, color = 'orange')


        #calculate the mid index of 
        ball_x = 10+9.5*math.sin(theta)
        ball_y = 10.5-9.5*math. cos(theta)
        mid = (projection(theta, phi, 10.0, 10.5)+projection(theta, phi, ball_x, ball_y))/2
        shade = max(3,abs(projection(theta, phi, 10.0, 10.5)-projection(theta, phi, ball_x, ball_y)))

        shadow = plt.Polygon(([mid - shade/2.0, -0.5],[mid + shade/2.0, -0.5]), color = 'black', linewidth = 3)
        
        ax = plt.gca()
        ax.add_artist(gun)
        ax.add_artist(ball)
        ax.add_artist(sun)
        ax.add_artist(shadow)
        ax.set_xlim((0, 20))
        ax.set_ylim((-1, 21))
        new=pd.DataFrame({
                  'i':(i-scale[0][0])/(scale[0][1]-0),
                  'j':(j-scale[1][0])/(scale[1][1]-0),
                  'shade':(shade-scale[2][0])/(scale[2][1]-0),
                  'mid':(mid-scale[2][0])/(scale[2][1]-0)
                  },
                  
                 index=[1])
        empty=empty.append(new,ignore_index=True)
        plt.axis('off')
        if count == 5:
          plt.savefig('/home/dyr/fd/DT-GNF_pendulum/causal_data/pendulum/test/a_' + str(round(float(i),4)) + '_' + str(round(float(j),4)) + '_' + str(round(float(shade),4)) + '_' + str(round(float(mid),4)) +'.png',dpi=96)
          count = 0
        else:
          plt.savefig('/home/dyr/fd/DT-GNF_pendulum/causal_data/pendulum/train/a_' + str(round(float(i),4)) + '_' + str(round(float(j),4)) + '_' + str(round(float(shade),4)) + '_' + str(round(float(mid),4)) +'.png',dpi=96)
        plt.clf()
        count += 1



