import os, glob
import pandas as pd
import numpy as np
import shapely.geometry as sgeom
import matplotlib.pyplot as plt
from matplotlib.image import imread
from matplotlib.animation import FuncAnimation
import matplotlib.lines as mlines
import cartopy.crs as ccrs
import cartopy.feature as cfeat
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.io.shapereader as shpreader
import cartopy.io.img_tiles as cimgt
from PIL import Image
import warnings 
warnings.filterwarnings('ignore')
df = pd.read_csv('./2018typhoon.csv')

def get_color(level):
  global color
  if level == '热带低压' or level == '热带扰动':
    color='#FFFF00'
  elif level == '热带风暴':
    color='#6495ED'
  elif level == '强热带风暴':
    color='#3CB371'
  elif level == '台风':
    color='#FFA500'
  elif level == '强台风':
    color='#FF00FF'
  elif level == '超强台风':
    color='#DC143C'
  return color

def create_map(title, extent):
  fig = plt.figure(figsize=(12, 8))
  ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
  url = 'http://map1c.vis.earthdata.nasa.gov/wmts-geo/wmts.cgi'
  layer = 'BlueMarble_ShadedRelief'
  ax.add_wmts(url, layer)
  ax.set_extent(extent,crs=ccrs.PlateCarree())

  gl = ax.gridlines(draw_labels=False, linewidth=1, color='k', alpha=0.5, linestyle='--')
  gl.xlabels_top = gl.ylabels_right = False 
  ax.set_xticks(np.arange(extent[0], extent[1]+5, 5))
  ax.set_yticks(np.arange(extent[2], extent[3]+5, 5))
  ax.xaxis.set_major_formatter(LongitudeFormatter())
  ax.xaxis.set_minor_locator(plt.MultipleLocator(1))
  ax.yaxis.set_major_formatter(LatitudeFormatter())
  ax.yaxis.set_minor_locator(plt.MultipleLocator(1))
  ax.tick_params(axis='both', labelsize=10, direction='out')

  a = mlines.Line2D([],[],color='#FFFF00',marker='o',markersize=7, label='TD',ls='')
  b = mlines.Line2D([],[],color='#6495ED', marker='o',markersize=7, label='TS',ls='')
  c = mlines.Line2D([],[],color='#3CB371', marker='o',markersize=7, label='STS',ls='')
  d = mlines.Line2D([],[],color='#FFA500', marker='o',markersize=7, label='TY',ls='')
  e = mlines.Line2D([],[],color='#FF00FF', marker='o',markersize=7, label='STY',ls='')
  f = mlines.Line2D([],[],color='#DC143C', marker='o',markersize=7, label='SSTY',ls='')
  ax.legend(handles=[a,b,c,d,e,f], numpoints=1, handletextpad=0, loc='upper left', shadow=True)
  plt.title(f'{title} Typhoon Track', fontsize=15)
  return ax

def draw_single(df):
  ax = create_map(df['名字'].iloc[0], [110, 135, 20, 45])
  for i in range(len(df)):
    ax.scatter(list(df['经度'])[i], list(df['纬度'])[i], marker='o', s=20, color=get_color(list(df['强度'])[i]))

  for i in range(len(df)-1):
    pointA = list(df['经度'])[i],list(df['纬度'])[i]
    pointB = list(df['经度'])[i+1],list(df['纬度'])[i+1]
    ax.add_geometries([sgeom.LineString([pointA, pointB])], color=get_color(list(df['强度'])[i+1]),crs=ccrs.PlateCarree())
  plt.savefig('./typhoon_one.png')
draw_single(df[df['编号']==1818])

def draw_multi(df):
  L = list(set(df['编号']))
  L.sort(key=list(df['编号']).index)
  ax = create_map('2018', [100, 180, 0, 45])
  for number in L:
    df1 = df[df['编号']==number]
    for i in range(len(df1)-1):
      pointA = list(df1['经度'])[i],list(df1['纬度'])[i]
      pointB = list(df1['经度'])[i+1],list(df1['纬度'])[i+1]
      ax.add_geometries([sgeom.LineString([pointA, pointB])], color=get_color(list(df1['强度'])[i+1]),crs=ccrs.PlateCarree())
  plt.savefig('./typhoon_multi.png')
draw_multi(df)

def draw_single_gif(df):
  for state in range(len(df.index))[:]:
    ax = create_map(f'{df["名字"].iloc[0]} {df["时间"].iloc[state]}', [110, 135, 20, 45])
    for i in range(len(df[:state])):
      ax.scatter(df['经度'].iloc[i], df['纬度'].iloc[i], marker='o', s=20, color=get_color(df['强度'].iloc[i]))
    for i in range(len(df[:state])-1):
      pointA = df['经度'].iloc[i],df['纬度'].iloc[i]
      pointB = df['经度'].iloc[i+1],df['纬度'].iloc[i+1]
      ax.add_geometries([sgeom.LineString([pointA, pointB])], color=get_color(df['强度'].iloc[i+1]),crs=ccrs.PlateCarree())
    print(f'正在绘制第{state}张轨迹图')
    plt.savefig(f'./{df["名字"].iloc[0]}{str(state).zfill(3)}.png', bbox_inches='tight')
  # 将图片拼接成动画
  imgFiles = list(glob.glob(f'./{df["名字"].iloc[0]}*.png'))
  images = [Image.open(fn) for fn in imgFiles]
  im = images[0]
  filename = f'./track_{df["名字"].iloc[0]}.gif'
  im.save(fp=filename, format='gif', save_all=True, append_images=images[1:], duration=500)
draw_single_gif(df[df['编号']==1818])