"""
用于刮取翼型轮廓的脚本http://airfoiltools.com/具有属性

Re=200000，Ncrit=9。

假设：

-外形定义从TE/SS开始，到TE/PS结束，并且是连续的（Selig）

-通过XFoil一致运行外形模拟

运行:
 > python3 airfoil_scrape.py --points
"""

################################################################################
# %% 引入包
################################################################################

import os
import argparse
import requests
from bs4 import BeautifulSoup
import re
import numpy as np
import ray
import more_itertools as mit

################################################################################
# %% 命令行参数分析

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

parser = argparse.ArgumentParser()
parser.add_argument("POINTS", help="number of points around circumference")
parser.add_argument("CPUS", help="number of CPUs to use for ray")
parser.add_argument("--cosine", help="use cosine spacing to refine LE", action="store_true")
args = parser.parse_args()

################################################################################
# %% 常量

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

BASE_URL = "http://airfoiltools.com"
POINTS = int(args.POINTS)
CPUS = int(args.CPUS)

################################################################################
# %% 如果还不存在，创建文件夹
################################################################################

try:
    os.stat(f'01-prep-data/{POINTS}')
except:
    os.mkdir(f'01-prep-data/{POINTS}')

################################################################################
# %% 初始化
################################################################################

ray.init(num_cpus=CPUS)

################################################################################
# %% 翼型扫描

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

##### 抓斗概述页面

page = requests.get(BASE_URL+"/search/airfoils/index.html")
overview_soup = BeautifulSoup(page.text, 'html.parser')

##### 查找所有链接
links = overview_soup.find_all("a", href=re.compile("airfoil/details"))
print(f'Found {len(links)} airfoil definitions')
links = [link.get('href') for link in links]

##### 分成CPU大小的块

chunk_links = [list(c) for c in mit.divide(CPUS, links)]

##### DEF远程功能在此！
@ray.remote
def scrape(links, BASE_URL):

    ##### 初始批次阵列
    X_batch = np.zeros((0, POINTS, 2, 1), dtype=float)
    y_batch = np.zeros((0, 3), dtype=float)

    ##### 回路连接

    for link in links:

        ##### 抓斗配置文件页面
        page = requests.get(BASE_URL+link)
        airfoil_soup = BeautifulSoup(page.text, 'html.parser')

        ##### 查找包含配置文件的数据文件
        dat_link = airfoil_soup.find("a", href=re.compile("seligdat"))

        profile_page = requests.get(BASE_URL+dat_link.get('href'))
        profile_soup = BeautifulSoup(profile_page.text, 'html.parser')

        text = profile_soup.text.split('\n')
        profile_name = text[0]
        raw_profile = np.loadtxt(text[1:])

        fp1 = raw_profile[:, 0]
        fp2 = raw_profile[:, 1]

        delta = np.sqrt(np.diff(fp1)**2 + np.diff(fp2)**2)
        xp = delta.cumsum()
        xp = np.insert(xp, 0, 0)/xp[-1]

        if not args.cosine:
            x = np.linspace(0, 1, POINTS)
        elif args.cosine:
            x = np.linspace(0,1, POINTS) + 0.1*np.sin(np.linspace(0, 2*np.pi, POINTS))

        base_profile = np.array([np.interp(x, xp, fp1), np.interp(x, xp, fp2)]).T

        xmin = base_profile[:, 0].min()
        xmax = base_profile[:, 0].max()
        base_profile[:, 0] = (base_profile[:, 0]-xmin)/(xmax-xmin)
        base_profile[:, 1] = (base_profile[:, 1])/(xmax-xmin)

        area = 0.0
        for i in range(len(base_profile)-1):
            area += 0.5*(base_profile[i, 0]-base_profile[i+1, 0])*(base_profile[i, 1]+base_profile[i+1, 1])

        data_link = airfoil_soup.find_all("a", href=re.compile("^/polar/details.*200000$"))

        if data_link and area <= 0.2:
            data_page = requests.get(BASE_URL+data_link[0].get('href'))
            data_soup = BeautifulSoup(data_page.text, 'html.parser')

            data = []
            table = data_soup.find('table', attrs={'class':'tabdata'})
            rows = table.find_all('tr')
            for row in rows[1:]:
                cols = row.find_all('td')
                cols = [ele.text.strip() for ele in cols]
                data.append([ele for ele in cols if ele])
            data = np.array(data).astype(float)[:, 0:3]

            X = np.zeros((len(data), POINTS, 2, 1), dtype=float)
            y = np.zeros((len(data), 3), dtype=float)

            for sample in range(len(data)):

                alpha = -data[sample, 0]*np.pi/180.0

                shift_profile = base_profile.copy()
                shift_profile[:, 0] = shift_profile[:, 0] - 0.5

                profile_x = shift_profile[:, 0]*np.cos(alpha) - shift_profile[:, 1]*np.sin(alpha)
                profile_y = shift_profile[:, 0]*np.sin(alpha) + shift_profile[:, 1]*np.cos(alpha)

                profile_x = profile_x + 0.5
                profile = np.concatenate((profile_x.reshape(-1, 1), profile_y.reshape(-1, 1)), axis=1)

                X[sample, :, :, 0] = profile
                y[sample, 0:2] = data[sample, 1:3]
                y[sample, 2] = area


            X_batch = np.concatenate((X_batch, X), axis=0)
            y_batch = np.concatenate((y_batch, y), axis=0)

    return X_batch, y_batch

################################################################################
# %% 并行执行
################################################################################

object_id = [scrape.remote(chunk_links[batch], BASE_URL) for batch in range(CPUS)]
out = ray.get(object_id)

################################################################################
# %% 关闭并保存文件
################################################################################

for i, batch in enumerate(out):

    ##### 获取本地数据
    X_batch = batch[0]
    y_batch = batch[1]

    idx = np.random.permutation(len(X_batch))
    X_batch = X_batch[idx]
    y_batch = y_batch[idx]

    ##### 保存
    np.save(f'01-prep-data/{POINTS}/X_{i:03d}.npy', X_batch)
    np.save(f'01-prep-data/{POINTS}/y_{i:03d}.npy', y_batch)