"""
CCT 建模优化代码
束线

作者：赵润晓
日期：2021年5月1日
"""

import multiprocessing  # since v0.1.1 多线程计算
import time  # since v0.1.1 统计计算时长
from typing import Callable, Dict, Generic, Iterable, List, NoReturn, Optional, Tuple, TypeVar, Union
import matplotlib.pyplot as plt
import math
import random  # since v0.1.1 随机数
import sys
import os  # since v0.1.1 查看CPU核心数
import numpy
from scipy.integrate import solve_ivp  # since v0.1.1 ODE45
import warnings  # since v0.1.1 提醒方法过时
from packages.point import *
from packages.constants import *
from packages.base_utils import BaseUtils
from packages.local_coordinate_system import LocalCoordinateSystem
from packages.line2s import *
from packages.trajectory import Trajectory
from packages.particles import *
from packages.magnets import *
from packages.cct import CCT

class Beamline(Line2, Magnet, ApertureObject):
    def __init__(self, trajectory: Trajectory) -> None:
        """
        不要直接调用构造器
        请使用 set_start_point
        """
        self.magnets: List[Magnet] = []
        self.trajectory: Trajectory = trajectory

        # 2021年3月18日 新增，表示元件。List 中每个元素表示一个元件
        # 元件由三部分组成，位置、元件自身、长度
        # 其中位置表示沿着 Beamline 的长度
        # 元件自身，使用 None 表示漂移段。
        self.elements: List[Tuple[float, Magnet, float]] = []

    def magnetic_field_at(self, point: P3) -> P3:
        """
        返回 Beamline 在全局坐标系点 P3 处产生的磁场
        """
        b: P3 = P3.zeros()
        for m in self.magnets:
            b += m.magnetic_field_at(point)
        return b

    # from Magnet
    def magnetic_field_along(
            self,
            line2: Optional[Line2] = None,
            p2_t0_p3: Callable[[P2], P3] = lambda p2: P3(p2.x, p2.y, 0.0),
            step: float = 1 * MM,
    ) -> List[ValueWithDistance[P3]]:
        """
        计算本对象在二维曲线 line2 上的磁场分布(line2 为 None 时，默认为 self.trajectory)
        p2_t0_p3 是一个函数，用于把 line2 上的二维点转为三维，默认转为 z=0 的三维点
        step 表示 line2 分段长度
        -------
        """
        if line2 is None:
            line2 = self.trajectory

        return super(Beamline, self).magnetic_field_along(
            line2=line2, p2_t0_p3=p2_t0_p3, step=step
        )

    def magnetic_field_bz_along(
            self,
            line2: Optional[Line2] = None,
            p2_t0_p3: Callable[[P2], P3] = lambda p2: P3(p2.x, p2.y, 0.0),
            step: float = 1 * MM,
    ) -> List[P2]:
        """
        计算本对象在二维曲线 line (line2 为 None 时，默认为 self.trajectory)上的磁场 Z 方向分量的分布
        因为磁铁一般放置在 XY 平面，所以 Bz 一般可以看作自然坐标系下 By，也就是二级场大小
        p2_t0_p3 是一个函数，用于把 line2 上的二维点转为三维，默认转为 z=0 的三维点
        step 表示 line2 分段长度

        返回 P2 的数组，P2 中 x 表示曲线 line2 上距离 s，y 表示前述距离对应的点的磁场 bz
        """
        if line2 is None:
            line2 = self.trajectory

        return super(Beamline, self).magnetic_field_bz_along(
            line2=line2, p2_t0_p3=p2_t0_p3, step=step
        )

    def graident_field_along(
            self,
            line2: Optional[Line2] = None,
            good_field_area_width: float = 10 * MM,
            step: float = 1 * MM,
            point_number: int = 4,
    ) -> List[P2]:
        """
        计算本对象在二维曲线 line2 (line2 为 None 时，默认为 self.trajectory)上的磁场梯度的分布
        每一点的梯度，采用这点水平垂线上 Bz 的多项式拟合得到
        good_field_area_width：水平垂线的长度，注意应小于等于好场区范围
        step：line2 上取点间距
        point_number：水平垂线上取点数目，越多则拟合越精确
        """
        if line2 is None:
            line2 = self.trajectory

        return super(Beamline, self).graident_field_along(
            line2=line2, good_field_area_width=good_field_area_width, step=step, point_number=point_number
        )

    def second_graident_field_along(
            self,
            line2: Optional[Line2] = None,
            good_field_area_width: float = 10 * MM,
            step: float = 1 * MM,
            point_number: int = 4,
    ) -> List[P2]:
        """
        计算本对象在二维曲线 line2 (line2 为 None 时，默认为 self.trajectory)上的磁场二阶梯度的分布（六极场）
        每一点的梯度，采用这点水平垂线上 Bz 的多项式拟合得到
        good_field_area_width：水平垂线的长度，注意应小于等于好场区范围
        step：line2 上取点间距
        point_number：水平垂线上取点数目，越多则拟合越精确
        """
        if line2 is None:
            line2 = self.trajectory

        return super(Beamline, self).second_graident_field_along(
            line2=line2, good_field_area_width=good_field_area_width, step=step, point_number=point_number
        )

    def track_ideal_particle(
            self,
            kinetic_MeV: float,
            s: float = 0.0,
            length: Optional[float] = None,
            footstep: float = 5 * MM,
    ) -> List[P3]:
        """
        束流跟踪，运行一个理想粒子，返回轨迹
        kinetic_MeV 粒子动能，单位 MeV
        s 起点位置
        length 粒子运行长度，默认运动到束线尾部
        footstep 粒子运动步长
        """
        if length is None:
            length = self.trajectory.get_length() - s
        ip = ParticleFactory.create_proton_along(
            self.trajectory, s, kinetic_MeV)
        return ParticleRunner.run_get_trajectory(ip, self, length, footstep)

    def track_phase_space_particle(
        self,
        x_mm: float,
        xp_mrad: float,
        y_mm: float,
        yp_mrad,
        delta: float,
        kinetic_MeV: float,
        s: float = 0.0,
        length: Optional[float] = None,
        footstep: float = 10 * MM,
    ) -> List[ValueWithDistance[PhaseSpaceParticle]]:
        """
        运行一个相空间粒子
        x_mm 相空间坐标 x，单位 mm
        xp_mrad 相空间坐标 xp，单位 mrad
        y_mm 相空间坐标 y，单位 mm
        yp_mrad 相空间坐标 yp，单位 mrad
        delta 动量分散
        kinetic_MeV 正则动能，单位 MeV
        s 在束线上的起点，默认 0.0
        length 运动长度，如果为空则运行到束线尾
        footstep 运动步长，默认 10*MM

        返回值是一个 List[ValueWithDistance[PhaseSpaceParticle]]
        即一个数组，数组元素是 ValueWithDistance
        即对应运动位置的粒子的相空间坐标信息
        """
        if length is None:
            length = self.trajectory.get_length() - s
        pp = PhaseSpaceParticle(
            x=x_mm * MM,
            xp=xp_mrad * MM,
            y=y_mm * MM,
            yp=yp_mrad * MM,
            z=0.0,
            delta=delta
        )
        # ip, distence = 0.0
        ip = ParticleFactory.create_proton_along(
            self.trajectory, s, kinetic_MeV)
        # to rp, distence = 0.0
        rp = ParticleFactory.create_from_phase_space_particle(
            ideal_particle=ip,
            coordinate_system=ip.get_natural_coordinate_system(),
            phase_space_particle=pp
        )
        # run all info, distence from 0.0
        all_info = ParticleRunner.run_get_all_info(
            p=rp,
            m=self,
            length=length,
            footstep=footstep
        )
        # for cp
        ret: List[ValueWithDistance[PhaseSpaceParticle]] = []
        for cp in all_info:
            d = cp.distance  # , distence from 0.0
            cip = ParticleFactory.create_proton_along(
                self.trajectory, d + s, kinetic_MeV)  # 所以这里是 d + s
            cpp = PhaseSpaceParticle.create_from_running_particle(
                ideal_particle=cip,
                coordinate_system=cip.get_natural_coordinate_system(),
                running_particle=cp
            )
            ret.append(ValueWithDistance(
                value=cpp, distance=d
            ))

        return ret

    def track_phase_ellipse(
            self,
            x_sigma_mm: float,
            xp_sigma_mrad: float,
            y_sigma_mm: float,
            yp_sigma_mrad,
            delta: float,
            particle_number: int,
            kinetic_MeV: float,
            s: float = 0.0,
            length: Optional[float] = None,
            footstep: float = 10 * MM,
            concurrency_level: int = 1,
            report: bool = True
    ) -> Tuple[List[P2], List[P2]]:
        """
        束流跟踪，运行两个相椭圆边界上的粒子，
        返回一个长度 2 的元组，表示相空间 x-xp 平面和 y-yp 平面上粒子投影（单位 mm / mrad）
        两个相椭圆，一个位于 xxp 平面，参数为 σx 和 σxp ，动量分散为 delta
        另一个位于 xxp 平面，参数为 σx 和 σxp ，动量分散为 delta
        x_sigma_mm σx 单位 mm
        xp_sigma_mrad σxp 单位 mrad
        y_sigma_mm σy 单位 mm
        yp_sigma_mrad σyp 单位 mrad
        delta 动量分散 单位 1
        particle_number 粒子数目
        kinetic_MeV 动能 单位 MeV
        s 起点位置
        length 粒子运行长度，默认运行到束线尾部
        footstep 粒子运动步长
        concurrency_level 并发等级（使用多少个核心进行粒子跟踪）
        report 是否打印并行任务计划
        """
        if length is None:
            length = self.trajectory.get_length() - s
        ip_start = ParticleFactory.create_proton_along(
            self.trajectory, s, kinetic_MeV)
        ip_end = ParticleFactory.create_proton_along(
            self.trajectory, s + length, kinetic_MeV
        )

        pp_x = PhaseSpaceParticle.phase_space_particles_along_positive_ellipse_in_xxp_plane(
            xMax=x_sigma_mm * MM,
            xpMax=xp_sigma_mrad * MRAD,
            delta=delta,
            number=particle_number,
        )

        pp_y = PhaseSpaceParticle.phase_space_particles_along_positive_ellipse_in_yyp_plane(
            yMax=y_sigma_mm * MM,
            ypMax=yp_sigma_mrad * MRAD,
            delta=delta,
            number=particle_number,
        )

        rp_x = ParticleFactory.create_from_phase_space_particles(
            ideal_particle=ip_start,
            coordinate_system=ip_start.get_natural_coordinate_system(),
            phase_space_particles=pp_x,
        )

        rp_y = ParticleFactory.create_from_phase_space_particles(
            ideal_particle=ip_start,
            coordinate_system=ip_start.get_natural_coordinate_system(),
            phase_space_particles=pp_y,
        )

        # run
        # refactor v0.1.1 合并计算
        ParticleRunner.run_only(
            p=rp_x + rp_y, m=self, length=length, footstep=footstep, concurrency_level=concurrency_level,
            report=report
        )

        pp_x_end = PhaseSpaceParticle.create_from_running_particles(
            ideal_particle=ip_end,
            coordinate_system=ip_end.get_natural_coordinate_system(),
            running_particles=rp_x,
        )

        pp_y_end = PhaseSpaceParticle.create_from_running_particles(
            ideal_particle=ip_end,
            coordinate_system=ip_end.get_natural_coordinate_system(),
            running_particles=rp_y,
        )

        xs = [pp.project_to_xxp_plane() / MM for pp in pp_x_end]
        ys = [pp.project_to_yyp_plane() / MM for pp in pp_y_end]

        s: BaseUtils.Statistic = BaseUtils.Statistic()

        print(
            f"delta={delta}," +
            f"avg_size_x={s.clear().add_all(P2.extract(xs)[0]).half_width()}mm," +
            f"avg_size_y={s.clear().add_all(P2.extract(ys)[0]).half_width()}mm"
        )

        return (xs, ys)

    # from ApertureObject
    def is_out_of_aperture(self, point: P3) -> bool:
        """
        判断点 point 是否超出 Beamline 的任意一个元件的孔径
        只有当粒子轴向投影在元件内部时，才会进行判断，
        否则即时粒子距离轴线很远，也认为粒子没有超出孔径，
        这是因为粒子不在元件内时，很可能处于另一个大孔径元件中，这样会造成误判。

        注意：这个函数的效率极低！
        """
        for m in self.magnets:
            if isinstance(m, ApertureObject) and m.is_out_of_aperture(point):
                print(f"beamline在{m}位置超出孔径")
                return True

        return False

    def trace_is_out_of_aperture(
            self, trace_with_distance: List[ValueWithDistance[P3]]
    ) -> bool:
        """
        判断一条粒子轨迹是否超出孔径

        注意：这个函数的效率极低！
        """
        for pd in trace_with_distance:
            if self.is_out_of_aperture(pd.value):
                return True

        return False

    def get_length(self) -> float:
        """
        获得 Beamline 的长度
        """
        return self.trajectory.get_length()

    def point_at(self, s: float) -> P2:
        """
        获得 Beamline s 位置处的点 (x,y)
        -------

        """
        return self.trajectory.point_at(s)

    def direct_at(self, s: float) -> P2:
        """
        获得 Beamline s 位置处的方向
        """
        return self.trajectory.direct_at(s)

    class __BeamlineBuilder:
        """
        构建 Beamline 的中间产物
        """

        def __init__(self, start_point: P2) -> None:
            self.start_point = start_point

        def first_drift(self, direct: P2 = P2.x_direct(), length: float = 1.0) -> "Beamline":
            """
            为 Beamline 添加第一个 drift
            正如 Trajectory 的第一个曲线段必须是是直线一样
            Beamline 中第一个元件必须是 drift
            """
            bl = Beamline(
                Trajectory.set_start_point(self.start_point).first_line(
                    direct=direct, length=length
                )
            )
            bl.elements.append((0, None, length))
            return bl

    @staticmethod
    # -> "Beamline.__BeamlineBuilder"
    def set_start_point(start_point: P2 = P2.origin()):
        """
        设置束线起点
        """
        return Beamline.__BeamlineBuilder(start_point)

    def append_drift(self, length: float) -> "Beamline":
        """
        尾加漂移段
        length 漂移段长度
        """
        old_len = self.trajectory.get_length()
        self.trajectory.add_strait_line(length=length)
        self.elements.append((old_len, None, length))

        return self

    def append_straight_dipole_magnet(
            self,
            magnetic_field: float,
            length: float,
            aperture_radius: float,
            # field_direct: P2 = P2.y_direct()
    ) -> "Beamline":
        """
        尾加直线二极铁
        """
        old_length = self.trajectory.get_length()
        self.trajectory.add_strait_line(length=length)

        lum = LocalUniformMagnet.create_local_uniform_magnet_along(
            trajectory=self.trajectory,
            s=old_length,
            length=length,
            magnetic_field=magnetic_field,
            aperture_radius=aperture_radius,
        )

        self.magnets.append(lum)
        self.elements.append((old_length, lum, length))

        return self

    def append_qs(
            self,
            length: float,
            gradient: float,
            second_gradient: float,
            aperture_radius: float,
    ) -> "Beamline":
        """
        尾加 QS 磁铁

        length: float QS 磁铁长度
        gradient: float 梯度 T/m
        second_gradient: float 二阶梯度（六极场） T/m^2
        aperture_radius: float 半孔径 单位 m
        """
        old_length = self.trajectory.get_length()
        self.trajectory.add_strait_line(length=length)

        qs = QS.create_qs_along(
            trajectory=self.trajectory,
            s=old_length,
            length=length,
            gradient=gradient,
            second_gradient=second_gradient,
            aperture_radius=aperture_radius,
        )

        self.magnets.append(qs)
        self.elements.append((old_length, qs, length))

        return self

    def append_q(
            self,
            length: float,
            gradient: float,
            aperture_radius: float,
    ) -> "Beamline":
        """
        尾加 Q 磁铁

        length: float QS 磁铁长度
        gradient: float 梯度 T/m
        aperture_radius: float 半孔径 单位 m
        """
        old_length = self.trajectory.get_length()
        self.trajectory.add_strait_line(length=length)

        q = Q.create_q_along(
            trajectory=self.trajectory,
            s=old_length,
            length=length,
            gradient=gradient,
            aperture_radius=aperture_radius,
        )

        self.magnets.append(q)
        self.elements.append((old_length, q, length))

        return self

    def append_dipole_cct(
            self,
            big_r: float,
            small_r_inner: float,
            small_r_outer: float,
            bending_angle: float,
            tilt_angles: List[float],
            winding_number: int,
            current: float,
            disperse_number_per_winding: int = 120,
    ) -> "Beamline":
        """
        尾加二极CCT

        big_r: float 偏转半径
        small_r_inner: float 内层半孔径
        small_r_outer: float 外层半孔径
        bending_angle: float 偏转角度（正数表示逆时针、负数表示顺时针）
        tilt_angles: List[float] 各极倾斜角
        winding_number: int 匝数
        current: float 电流
        disperse_number_per_winding: int 每匝分段数目，越大计算越精确
        """
        old_length = self.trajectory.get_length()
        cct_length = big_r * abs(BaseUtils.angle_to_radian(bending_angle))
        self.trajectory.add_arc_line(
            radius=big_r, clockwise=bending_angle < 0, angle_deg=abs(bending_angle)
        )

        cct_inner = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=small_r_inner,
            bending_angle=abs(bending_angle),
            tilt_angles=tilt_angles,
            winding_number=winding_number,
            current=current,
            starting_point_in_ksi_phi_coordinate=P2.origin(),
            end_point_in_ksi_phi_coordinate=P2(
                2 * math.pi * winding_number,
                BaseUtils.angle_to_radian(bending_angle),
            ),
            disperse_number_per_winding=disperse_number_per_winding,
        )
        self.magnets.append(cct_inner)
        self.elements.append((old_length, cct_inner, cct_length))

        cct_outer = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=small_r_outer,
            bending_angle=abs(bending_angle),
            tilt_angles=BaseUtils.list_multiply(tilt_angles, -1),
            winding_number=winding_number,
            current=current,
            starting_point_in_ksi_phi_coordinate=P2.origin(),
            end_point_in_ksi_phi_coordinate=P2(
                -2 * math.pi * winding_number,
                BaseUtils.angle_to_radian(bending_angle),
            ),
            disperse_number_per_winding=disperse_number_per_winding,
        )
        self.magnets.append(cct_outer)
        self.elements.append((old_length, cct_outer, cct_length))

        return self

    def append_agcct(
            self,
            big_r: float,
            small_rs: List[float],
            bending_angles: List[float],
            tilt_angles: List[List[float]],
            winding_numbers: List[List[int]],
            currents: List[float],
            disperse_number_per_winding: int = 120,
    ) -> "Beamline":
        """
        尾加 agcct
        本质是两层二极 CCT 和两层交变四极 CCT

        big_r: float 偏转半径，单位 m
        small_rs: List[float] 各层 CCT 的孔径，一共四层，从大到小排列。分别是二极CCT外层、内层，四极CCT外层、内层
        bending_angles: List[float] 交变四极 CCT 每个 part 的偏转半径（正数表示逆时针、负数表示顺时针），要么全正数，要么全负数。不需要传入二极 CCT 偏转半径，因为它就是 sum(bending_angles)
        tilt_angles: List[List[float]] 二极 CCT 和四极 CCT 的倾斜角，典型值 [[30],[90,30]]，只有两个元素的二维数组
        winding_numbers: List[List[int]], 二极 CCT 和四极 CCT 的匝数，典型值 [[128],[21,50,50]] 表示二极 CCT 128匝，四极交变 CCT 为 21、50、50 匝
        currents: List[float] 二极 CCT 和四极 CCT 的电流，典型值 [8000,9000]
        disperse_number_per_winding: int 每匝分段数目，越大计算越精确

        添加 CCT 的顺序为：
        外层二极 CCT
        内层二极 CCT
        part1 四极 CCT 内层
        part1 四极 CCT 外层
        part2 四极 CCT 内层
        part2 四极 CCT 外层
        ... ...
        """
        if len(small_rs) != 4:
            raise ValueError(
                f"small_rs({small_rs})，长度应为4，分别是二极CCT外层、内层，四极CCT外层、内层")
        if not BaseUtils.is_sorted(small_rs[::-1]):
            raise ValueError(
                f"small_rs({small_rs})，应从大到小排列，分别是二极CCT外层、内层，四极CCT外层、内层")

        total_bending_angle = sum(bending_angles)
        old_length = self.trajectory.get_length()
        cct_length = big_r * \
            abs(BaseUtils.angle_to_radian(total_bending_angle))
        self.trajectory.add_arc_line(
            radius=big_r,
            clockwise=total_bending_angle < 0,
            angle_deg=abs(total_bending_angle),
        )

        # 构建二极 CCT 外层
        cct2_outer = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=small_rs[0],
            bending_angle=abs(total_bending_angle),
            tilt_angles=BaseUtils.list_multiply(tilt_angles[0], -1),
            winding_number=winding_numbers[0][0],
            current=currents[0],
            starting_point_in_ksi_phi_coordinate=P2.origin(),
            end_point_in_ksi_phi_coordinate=P2(
                -2 * math.pi * winding_numbers[0][0],
                BaseUtils.angle_to_radian(total_bending_angle),
            ),
            disperse_number_per_winding=disperse_number_per_winding,
        )
        self.magnets.append(cct2_outer)
        self.elements.append((old_length, cct2_outer, cct_length))

        # 构建二极 CCT 内层
        cct2_innter = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=small_rs[1],
            bending_angle=abs(total_bending_angle),
            tilt_angles=tilt_angles[0],
            winding_number=winding_numbers[0][0],
            current=currents[0],
            starting_point_in_ksi_phi_coordinate=P2.origin(),
            end_point_in_ksi_phi_coordinate=P2(
                2 * math.pi * winding_numbers[0][0],
                BaseUtils.angle_to_radian(total_bending_angle),
            ),
            disperse_number_per_winding=disperse_number_per_winding,
        )
        self.magnets.append(cct2_innter)
        self.elements.append((old_length, cct2_innter, cct_length))

        # 构建内外侧四极交变 CCT
        # 提取参数
        agcct_small_r_out = small_rs[2]
        agcct_small_r_in = small_rs[3]
        agcct_winding_nums: List[int] = winding_numbers[1]
        agcct_bending_angles: List[float] = bending_angles
        agcct_bending_angles_rad: List[float] = BaseUtils.angle_to_radian(
            agcct_bending_angles
        )
        agcct_tilt_angles: List[float] = tilt_angles[1]
        agcct_current: float = currents[1]

        # 构建 part1
        agcct_index = 0
        agcct_start_in = P2.origin()
        agcct_start_out = P2.origin()
        agcct_end_in = P2(
            ((-1.0) ** agcct_index) * 2 * math.pi *
            agcct_winding_nums[agcct_index],
            agcct_bending_angles_rad[agcct_index],
        )
        agcct_end_out = P2(
            ((-1.0) ** (agcct_index + 1))
            * 2
            * math.pi
            * agcct_winding_nums[agcct_index],
            agcct_bending_angles_rad[agcct_index],
        )
        agcct_part1_inner = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=agcct_small_r_in,
            bending_angle=abs(agcct_bending_angles[agcct_index]),
            tilt_angles=BaseUtils.list_multiply(agcct_tilt_angles, -1),
            winding_number=agcct_winding_nums[agcct_index],
            current=agcct_current,
            starting_point_in_ksi_phi_coordinate=agcct_start_in,
            end_point_in_ksi_phi_coordinate=agcct_end_in,
            disperse_number_per_winding=disperse_number_per_winding,
        )
        agcct_part1_length = big_r * \
            BaseUtils.angle_to_radian(abs(agcct_bending_angles[agcct_index]))
        self.magnets.append(agcct_part1_inner)
        self.elements.append(
            (old_length, agcct_part1_inner, agcct_part1_length))

        agcct_part1_outer = CCT.create_cct_along(
            trajectory=self.trajectory,
            s=old_length,
            big_r=big_r,
            small_r=agcct_small_r_out,
            bending_angle=abs(agcct_bending_angles[agcct_index]),
            tilt_angles=agcct_tilt_angles,
            winding_number=agcct_winding_nums[agcct_index],
            current=agcct_current,
            starting_point_in_ksi_phi_coordinate=agcct_start_out,
            end_point_in_ksi_phi_coordinate=agcct_end_out,
            disperse_number_per_winding=disperse_number_per_winding,
        )
        self.magnets.append(agcct_part1_outer)
        self.elements.append(
            (old_length, agcct_part1_outer, agcct_part1_length))

        old_length_i = old_length + agcct_part1_length
        # 构建 part2 和之后的 part
        for ignore in range(len(agcct_bending_angles) - 1):
            agcct_index += 1
            agcct_start_in = agcct_end_in + P2(
                0,
                agcct_bending_angles_rad[agcct_index - 1]
                / agcct_winding_nums[agcct_index - 1],
            )
            agcct_start_out = agcct_end_out + P2(
                0,
                agcct_bending_angles_rad[agcct_index - 1]
                / agcct_winding_nums[agcct_index - 1],
            )
            agcct_end_in = agcct_start_in + P2(
                ((-1) ** agcct_index) * 2 * math.pi *
                agcct_winding_nums[agcct_index],
                agcct_bending_angles_rad[agcct_index],
            )
            agcct_end_out = agcct_start_out + P2(
                ((-1) ** (agcct_index + 1))
                * 2
                * math.pi
                * agcct_winding_nums[agcct_index],
                agcct_bending_angles_rad[agcct_index],
            )
            agcct_parti_inner = CCT.create_cct_along(
                trajectory=self.trajectory,
                s=old_length,
                big_r=big_r,
                small_r=agcct_small_r_in,
                bending_angle=abs(agcct_bending_angles[agcct_index]),
                tilt_angles=BaseUtils.list_multiply(agcct_tilt_angles, -1),
                winding_number=agcct_winding_nums[agcct_index],
                current=agcct_current,
                starting_point_in_ksi_phi_coordinate=agcct_start_in,
                end_point_in_ksi_phi_coordinate=agcct_end_in,
                disperse_number_per_winding=disperse_number_per_winding,
            )
            agcct_parti_length = big_r * \
                BaseUtils.angle_to_radian(
                    abs(agcct_bending_angles[agcct_index]))
            self.magnets.append(agcct_parti_inner)
            self.elements.append(
                (old_length_i, agcct_parti_inner, agcct_parti_length))

            agcct_parti_outer = CCT.create_cct_along(
                trajectory=self.trajectory,
                s=old_length,
                big_r=big_r,
                small_r=agcct_small_r_out,
                bending_angle=abs(agcct_bending_angles[agcct_index]),
                tilt_angles=agcct_tilt_angles,
                winding_number=agcct_winding_nums[agcct_index],
                current=agcct_current,
                starting_point_in_ksi_phi_coordinate=agcct_start_out,
                end_point_in_ksi_phi_coordinate=agcct_end_out,
                disperse_number_per_winding=disperse_number_per_winding,
            )
            self.magnets.append(agcct_parti_outer)
            self.elements.append(
                (old_length_i, agcct_parti_outer, agcct_parti_length))

            old_length_i += agcct_parti_length

        return self

    def get_magnets(self) -> List[Magnet]:
        return self.magnets

    def get_trajectory(self) -> Trajectory:
        return self.trajectory

    def __str__(self) -> str:
        return f"beamline(magnet_size={len(self.magnets)}, traj_len={self.trajectory.get_length()})"

    def __repr__(self) -> str:
        return self.__str__()


# -*- coding: UTF-8 -*-

from mpi4py import MPI
from sympy import pi, cos, sin
import pytest
import os

from sympde.calculus import grad, dot
from sympde.topology import ScalarFunctionSpace, VectorFunctionSpace
from sympde.topology import ProductSpace
from sympde.topology import element_of
from sympde.topology import NormalVector
from sympde.topology import Union
from sympde.topology import Domain
from sympde.expr     import BilinearForm, LinearForm, integral
from sympde.expr     import Norm
from sympde.expr     import find, EssentialBC

from psydac.fem.basic          import FemField
from psydac.api.discretization import discretize

# ... get the mesh directory
try:
    mesh_dir = os.environ['PSYDAC_MESH_DIR']

except:
    base_dir = os.path.dirname(os.path.realpath(__file__))
    base_dir = os.path.join(base_dir, '..', '..', '..')
    mesh_dir = os.path.join(base_dir, 'mesh')
# ...

#==============================================================================
def run_poisson_3d_dir(filename, solution, f, comm=None):

    # ... abstract model
    domain = Domain.from_file(filename)

    V = ScalarFunctionSpace('V', domain)

    x,y,z = domain.coordinates

    F = element_of(V, name='F')

    v = element_of(V, name='v')
    u = element_of(V, name='u')

    int_0 = lambda expr: integral(domain , expr)

    expr = dot(grad(v), grad(u))
    a = BilinearForm((v,u), int_0(expr))

    expr = f*v
    l = LinearForm(v, int_0(expr))

    error = F - solution
    l2norm = Norm(error, domain, kind='l2')
    h1norm = Norm(error, domain, kind='h1')

    bc = EssentialBC(u, 0, domain.boundary)
    equation = find(u, forall=v, lhs=a(u,v), rhs=l(v), bc=bc)
    # ...

    # ... create the computational domain from a topological domain
    domain_h = discretize(domain, filename=filename, comm=comm)
    # ...

    # ... discrete spaces
    Vh = discretize(V, domain_h)
    # ...

    # ... dsicretize the equation using Dirichlet bc
    equation_h = discretize(equation, domain_h, [Vh, Vh])
    # ...

    # ... discretize norms
    l2norm_h = discretize(l2norm, domain_h, Vh)
    h1norm_h = discretize(h1norm, domain_h, Vh)
    # ...

    # ... solve the discrete equation
    x = equation_h.solve()
    # ...

    # ...
    phi = FemField( Vh, x )
    # ...

    # ... compute norms
    l2_error = l2norm_h.assemble(F=phi)
    h1_error = h1norm_h.assemble(F=phi)
    # ...

    return l2_error, h1_error

#==============================================================================
def run_poisson_3d_dirneu(filename, solution, f, boundary, comm=None):

    assert( isinstance(boundary, (list, tuple)) )

    # ... abstract model
    domain = Domain.from_file(filename)

    V = ScalarFunctionSpace('V', domain)

    B_neumann = [domain.get_boundary(**kw) for kw in boundary]
    if len(B_neumann) == 1:
        B_neumann = B_neumann[0]

    else:
        B_neumann = Union(*B_neumann)

    x,y,z = domain.coordinates

    F = element_of(V, name='F')

    v = element_of(V, name='v')
    u = element_of(V, name='u')

    nn = NormalVector('nn')

    int_0 = lambda expr: integral(domain , expr)
    int_1 = lambda expr: integral(B_neumann , expr)

    expr = dot(grad(v), grad(u))
    a = BilinearForm((v,u), int_0(expr))

    expr = f*v
    l0 = LinearForm(v, int_0(expr))

    expr = v*dot(grad(solution), nn)
    l_B_neumann = LinearForm(v, int_1(expr))

    expr = l0(v) + l_B_neumann(v)
    l = LinearForm(v, expr)

    error = F-solution
    l2norm = Norm(error, domain, kind='l2')
    h1norm = Norm(error, domain, kind='h1')

    B_dirichlet = domain.boundary.complement(B_neumann)
    bc = EssentialBC(u, 0, B_dirichlet)

    equation = find(u, forall=v, lhs=a(u,v), rhs=l(v), bc=bc)
    # ...

    # ... create the computational domain from a topological domain
    domain_h = discretize(domain, filename=filename, comm=comm)
    # ...

    # ... discrete spaces
    Vh = discretize(V, domain_h)
    # ...

    # ... dsicretize the equation using Dirichlet bc
    equation_h = discretize(equation, domain_h, [Vh, Vh])
    # ...

    # ... discretize norms
    l2norm_h = discretize(l2norm, domain_h, Vh)
    h1norm_h = discretize(h1norm, domain_h, Vh)
    # ...

    # ... solve the discrete equation
    x = equation_h.solve()
    # ...

    # ...
    phi = FemField( Vh, x )
    # ...

    # ... compute norms
    l2_error = l2norm_h.assemble(F=phi)
    h1_error = h1norm_h.assemble(F=phi)
    # ...

    return l2_error, h1_error

#==============================================================================
def run_laplace_3d_neu(filename, solution, f, comm=None):

    # ... abstract model
    domain = Domain.from_file(filename)

    V = ScalarFunctionSpace('V', domain)

    B_neumann = domain.boundary

    x,y,z = domain.coordinates

    F = element_of(V, name='F')

    v = element_of(V, name='v')
    u = element_of(V, name='u')

    nn = NormalVector('nn')

    int_0 = lambda expr: integral(domain , expr)
    int_1 = lambda expr: integral(B_neumann , expr)

    expr = dot(grad(v), grad(u)) + v*u
    a = BilinearForm((v,u), int_0(expr))

    expr = f*v
    l0 = LinearForm(v, int_0(expr))

    expr = v*dot(grad(solution), nn)
    l_B_neumann = LinearForm(v, int_1(expr))

    expr = l0(v) + l_B_neumann(v)
    l = LinearForm(v, expr)

    error = F-solution
    l2norm = Norm(error, domain, kind='l2')
    h1norm = Norm(error, domain, kind='h1')

    equation = find(u, forall=v, lhs=a(u,v), rhs=l(v))
    # ...

    # ... create the computational domain from a topological domain
    domain_h = discretize(domain, filename=filename, comm=comm)
    # ...

    # ... discrete spaces
    Vh = discretize(V, domain_h)
    # ...

    # ... dsicretize the equation using Dirichlet bc
    equation_h = discretize(equation, domain_h, [Vh, Vh])
    # ...

    # ... discretize norms
    l2norm_h = discretize(l2norm, domain_h, Vh)
    h1norm_h = discretize(h1norm, domain_h, Vh)
    # ...

    # ... solve the discrete equation
    x = equation_h.solve()
    # ...

    # ...
    phi = FemField( Vh, x )
    # ...

    # ... compute norms
    l2_error = l2norm_h.assemble(F=phi)
    h1_error = h1norm_h.assemble(F=phi)
    # ...

    return l2_error, h1_error

###############################################################################
#            SERIAL TESTS
###############################################################################

#==============================================================================
def test_api_poisson_3d_dir_collela():

    filename = os.path.join(mesh_dir, 'collela_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(pi*x)*sin(pi*y)*sin(pi*z)
    f        = 3*pi**2*sin(pi*x)*sin(pi*y)*sin(pi*z)

    l2_error, h1_error = run_poisson_3d_dir(filename, solution, f)

    expected_l2_error =  0.15687494944868827
    expected_h1_error =  1.518006054794389

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_poisson_3d_dirneu_identity_2():
    filename = os.path.join(mesh_dir, 'identity_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(0.5*pi*x)*sin(pi*y)*sin(pi*z)
    f        = (9./4.)*pi**2*solution

    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f, [{'axis': 0, 'ext': 1}])

    expected_l2_error =  0.001438835012218704
    expected_h1_error =  0.03929404299152016

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_poisson_3d_dirneu_identity_13():
    filename = os.path.join(mesh_dir, 'identity_3d.h5')

    from sympy.abc import x,y,z

    solution = cos(0.5*pi*x)*cos(0.5*pi*y)*sin(pi*z)
    f        = (3./2.)*pi**2*solution

    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
                                               [{'axis': 0, 'ext': -1},
                                                {'axis': 1, 'ext': -1}])

    expected_l2_error =  0.0010275451113313282
    expected_h1_error =  0.027938446826372126

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_poisson_3d_dirneu_identity_24():
    filename = os.path.join(mesh_dir, 'identity_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(0.5*pi*x)*sin(0.5*pi*y)*sin(pi*z)
    f        = (3./2.)*pi**2*solution

    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
                                               [{'axis': 0, 'ext': 1},
                                                {'axis': 1, 'ext': 1}])

    expected_l2_error =  0.001027545111330973
    expected_h1_error =  0.027938446826371813

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_poisson_3d_dirneu_identity_123():
#    filename = os.path.join(mesh_dir, 'identity_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = cos(0.25*pi*x)*cos(0.5*pi*y)*sin(pi*z)
#    f        = (21./16.)*pi**2*solution
#
#    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
#                                               [{'axis': 0, 'ext': -1},
#                                                {'axis': 0, 'ext': 1},
#                                                {'axis': 1, 'ext': -1}])
#
#    expected_l2_error =  0.0013124098938804697
#    expected_h1_error =  0.035441679549890456
#
#    assert( abs(l2_error - expected_l2_error) < 1.e-7)
#    assert( abs(h1_error - expected_h1_error) < 1.e-7)

##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_poisson_3d_dirneu_identity_1235():
#    filename = os.path.join(mesh_dir, 'identity_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = cos(0.25*pi*x)*cos(0.5*pi*y)*cos(0.5*pi*z)
#    f        = (9./16.)*pi**2*solution
#
#    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
#                                               [{'axis': 0, 'ext': -1},
#                                                {'axis': 0, 'ext': 1},
#                                                {'axis': 1, 'ext': -1},
#                                                {'axis': 2, 'ext': -1}])
#
#    expected_l2_error =  0.00019677816039781896
#    expected_h1_error =  0.0058786142515790405
#
#    assert( abs(l2_error - expected_l2_error) < 1.e-7)
#    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_poisson_3d_dirneu_collela_2():
    filename = os.path.join(mesh_dir, 'collela_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(0.25*pi*(x+1.))*sin(pi*y)*sin(pi*z)
    f        = (33./16.)*pi**2*solution

    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
                                               [{'axis': 0, 'ext': 1}])

    expected_l2_error =  0.06091240085930318
    expected_h1_error =  0.6380043932563333

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_poisson_3d_dirneu_collela_13():
#    filename = os.path.join(mesh_dir, 'collela_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = sin(0.25*pi*(1.-x))*sin(0.25*pi*(1.-y))*sin(pi*z)
#    f        = (9./8.)*pi**2*solution
#
#    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
#                                               [{'axis': 0, 'ext': -1},
#                                                {'axis': 1, 'ext': -1}])
#
#    expected_l2_error =  0.03786854933218588
#    expected_h1_error =  0.38437667047918933
#
#    assert( abs(l2_error - expected_l2_error) < 1.e-7)
#    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_poisson_3d_dirneu_collela_24():
    filename = os.path.join(mesh_dir, 'collela_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(0.25*pi*(x+1.))*sin(0.25*pi*(y+1.))*sin(pi*z)
    f        = (9./8.)*pi**2*solution

    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
                                               [{'axis': 0, 'ext': 1},
                                                {'axis': 1, 'ext': 1}])

    expected_l2_error =  0.03793880183960465
    expected_h1_error =  0.38439642303250143

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_poisson_3d_dirneu_collela_123():
#    filename = os.path.join(mesh_dir, 'collela_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = cos(pi*x)*sin(0.25*pi*(1.-y))*sin(pi*z)
#    f        = (33./16.)*pi**2*solution
#
#    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
#                                               [{'axis': 0, 'ext': -1},
#                                                {'axis': 0, 'ext': 1},
#                                                {'axis': 1, 'ext': -1}])
#
#    expected_l2_error =  0.11963989196330076
#    expected_h1_error =  1.1267766354124575
#
#    assert( abs(l2_error - expected_l2_error) < 1.e-7)
#    assert( abs(h1_error - expected_h1_error) < 1.e-7)
#
##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_poisson_3d_dirneu_collela_1235():
#    filename = os.path.join(mesh_dir, 'collela_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = cos(pi*x)*sin(0.25*pi*(1.-y))*sin(0.25*pi*(1.-z))
#    f        = (9./8.)*pi**2*solution
#
#    l2_error, h1_error = run_poisson_3d_dirneu(filename, solution, f,
#                                               [{'axis': 0, 'ext': -1},
#                                                {'axis': 0, 'ext': 1},
#                                                {'axis': 1, 'ext': -1},
#                                                {'axis': 2, 'ext': -1}])
#
#    expected_l2_error =  0.13208728319093133
#    expected_h1_error =  0.9964934429086868

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
def test_api_laplace_3d_neu_identity():
    filename = os.path.join(mesh_dir, 'identity_3d.h5')

    from sympy.abc import x,y,z

    solution = cos(pi*x)*cos(pi*y)*cos(pi*z)
    f        = (3.*pi**2 + 1.)*solution

    l2_error, h1_error = run_laplace_3d_neu(filename, solution, f)

    expected_l2_error =  0.0016975430150953524
    expected_h1_error =  0.047009063231215

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

##==============================================================================
## TODO DEBUG, not working since merge with devel
#def test_api_laplace_3d_neu_collela():
#    filename = os.path.join(mesh_dir, 'collela_3d.h5')
#
#    from sympy.abc import x,y,z
#
#    solution = cos(pi*x)*cos(pi*y)*cos(pi*z)
#    f        = (3.*pi**2 + 1.)*solution
#
#    l2_error, h1_error = run_laplace_3d_neu(filename, solution, f)
#
#    expected_l2_error =  0.1768000505351402
#    expected_h1_error =  1.7036022067226382
#
#    assert( abs(l2_error - expected_l2_error) < 1.e-7)
#    assert( abs(h1_error - expected_h1_error) < 1.e-7)

###############################################################################
#            PARALLEL TESTS
###############################################################################

#==============================================================================
@pytest.mark.parallel
def test_api_poisson_3d_dir_collela():

    filename = os.path.join(mesh_dir, 'collela_3d.h5')

    from sympy.abc import x,y,z

    solution = sin(pi*x)*sin(pi*y)*sin(pi*z)
    f        = 3*pi**2*sin(pi*x)*sin(pi*y)*sin(pi*z)

    l2_error, h1_error = run_poisson_3d_dir(filename, solution, f,
                                            comm=MPI.COMM_WORLD)

    expected_l2_error =  0.15687494944868827
    expected_h1_error =  1.518006054794389

    assert( abs(l2_error - expected_l2_error) < 1.e-7)
    assert( abs(h1_error - expected_h1_error) < 1.e-7)

#==============================================================================
# CLEAN UP SYMPY NAMESPACE
#==============================================================================

def teardown_module():
    from sympy import cache
    cache.clear_cache()

def teardown_function():
    from sympy import cache
    cache.clear_cache()


<filename>modules/math-codes/modules/statistics-and-probability/src/skew_kurt-v1.py
########################################################
# <NAME> - drigols                              #
# Last update: 17/12/2021                              #
########################################################

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
import scipy.stats as stats

df = pd.DataFrame({
  'Name': ['Dan', 'Joann', 'Pedro', 'Rosie', 'Ethan', 'Vicky', 'Frederic'],
  'Salary':[50000, 54000, 50000, 189000, 55000, 40000, 59000],
  'Hours':[41, 40, 36, 30, 35, 39, 40],
  'Grade':[50, 50, 46, 95, 50, 5, 57]
})

# Cria uma lista que vai representar os labels do nosso DataFrame:
# - Salary;
# - Hours;
# - Grade.
numcols = ['Salary', 'Hours', 'Grade']

# O laço for vai passar por cada item na nosssa lista de labels(numcols) fazendo o seguinte:
# - Imprimindo a assimétria (skewness);
# - Imprimindo a curtose (kurtosis);
# - Pegando a densidade da label;
# - Criando um Histograma para a label;
# - Adicionando a densidade/linha de densidade no plot/histograma.
for col in numcols:
  print(df[col].name + ' skewness: ' + str(df[col].skew())) # Imprime a Assimetria do label/coluna no laço for.
  print(df[col].name + ' kurtosis: ' + str(df[col].kurt())) # Imprime a Curtose do label/coluna no laço for.
  density = stats.gaussian_kde(df[col]) # Pega a densidade do label/coluna no laço for.
  n, x, _ = plt.hist(df[col], histtype='step', density=True, bins=25) # Cria o plot do label/coluna no laço for.
  plt.plot(x, density(x)*6) # Cria a linha de densidade do label no laço for.
  plt.show()
  print('\n')


<reponame>samcom12/anuga_core
"""

Function which can be useful when setting quantities

"""

import copy
import os
import anuga.utilities.spatialInputUtil as su

def make_nearestNeighbour_quantity_function(
    quantity_xyValueIn,
    domain,
    threshold_distance=9.0e+100,
    background_value=9.0e+100,
    k_nearest_neighbours=1,
    method='average'
):
    """
    Function which makes another function, which can be used in set_quantity
    Idea: For every point x,y in the domain, we want to set a quantity based on
          the 'nearest-neighbours' from quantity_xyValue (a 3 column array with
          x,y,quantity-value),
          UNLESS the distance from x,y to the nearest-neighbour is >
            threshold_distance.
          In the latter case, we want to set the quantity value to
            'background_value'

          We need a function f(x,y) to do that. This routine makes the
          function, with the desired quantity_xyValue points,
          threshold_distance, and background_value
    INPUTS:
        @param quantity_xyValueIn -- A 3 column array with 'x,y, Value'
            defining the points used to set the new quantity values in
            georeferenced coordinates
        @param domain -- The ANUGA domain
        @param k_nearest_neighbors --- Number of nearest neighbours used in calculation
        @param threshold_distance -- Points greater than this distance from
            their nearest quantity_xyValue point are set to background_value
        @param background_value -- see 'threshold_distance'
        @param method -- Three methods; 'average' uses an inverse-distance-weighted average
               of the k nearest neighbours is used:
               'min' the minimum of the k nearest neighbours is used:
               'max' the maximum of the k nearest neighbours is used.
    OUTPUTS:
        A function f which can be passed to domain.set_quantity('myQuantity', f)
    """

    import scipy
    import scipy.interpolate
    import scipy.spatial

    if(len(quantity_xyValueIn.shape) > 1):
        quantity_xyValue = quantity_xyValueIn
    else:
        # Treat the single-point case
        quantity_xyValue = quantity_xyValueIn.reshape((1, 3))

    # Make a function which gives us the ROW-INDEX of the nearest xy point in
    # quantity_xyValue
    # quantity_xy_interpolator = scipy.interpolate.NearestNDInterpolator(
    #    quantity_xyValue[:,0:2],
    #    scipy.arange(len(quantity_xyValue[:,2])))

    # Make a function which returns k-nearest-neighbour indices + distances
    quantity_xy_interpolator = scipy.spatial.cKDTree(quantity_xyValue[:, 0:2])

    # Make a function of x,y which we can pass to domain.set_quantity
    def quant_NN_fun(x, y):
        """
        Function to assign quantity from the nearest point in quantity_xyValue,
        UNLESS the point is more than 'threshold_distance' away from the
        nearest point, in which case the background value is used
        """

        import scipy
        import scipy.interpolate
        import scipy.spatial
        import numpy as np

        x = np.asarray(x).reshape(1, -1)[0, :]
        y = np.asarray(y).reshape(1, -1)[0, :]

        # Since ANUGA stores x,y internally in non-georeferenced coordinates,
        # we adjust them here
        xll = domain.geo_reference.xllcorner
        yll = domain.geo_reference.yllcorner
        z = np.zeros(shape=(len(x), 2))
        z[:, 0] = x+xll
        z[:, 1] = y+yll

        # This will hold the quantity values
        quantity_output = x*0. + background_value
        # Compute the index of the nearest-neighbour in quantity_xyValue
        neighbour_data = quantity_xy_interpolator.query(z,
                                                        k=k_nearest_neighbours)

        # Next find indices with distance < threshold_distance
        if(k_nearest_neighbours == 1):
            dist_lt_thresh = neighbour_data[0] < threshold_distance
        else:
            dist_lt_thresh = neighbour_data[0][:, 0] < threshold_distance

        dist_lt_thresh = dist_lt_thresh.nonzero()[0]

        # Initialise output
        quantity_output = x*0 + background_value

        # Interpolate
        if len(dist_lt_thresh) > 0:
            if method == 'min':
                numerator = 9.0e+100
                for i in range(k_nearest_neighbours):
                    if(k_nearest_neighbours == 1):
                        distances = neighbour_data[0][dist_lt_thresh]
                        indices = neighbour_data[1][dist_lt_thresh]
                        values = quantity_xyValue[indices, 2]
                        numerator = np.minimum(numerator, values)
                    else:
                        distances = neighbour_data[0][dist_lt_thresh, i]
                        indices = neighbour_data[1][dist_lt_thresh, i]
                        values = quantity_xyValue[indices, 2]
                        numerator = np.minimum(numerator, values)
                quantity_output[dist_lt_thresh] = numerator
            elif method == 'max':
                numerator = -9.0e+100
                for i in range(k_nearest_neighbours):
                    if(k_nearest_neighbours == 1):
                        distances = neighbour_data[0][dist_lt_thresh]
                        indices = neighbour_data[1][dist_lt_thresh]
                        values = quantity_xyValue[indices, 2]
                        numerator = np.maximum(numerator, values)
                    else:
                        distances = neighbour_data[0][dist_lt_thresh, i]
                        indices = neighbour_data[1][dist_lt_thresh, i]
                        values = quantity_xyValue[indices, 2]
                        numerator = np.maximum(numerator, values)
                quantity_output[dist_lt_thresh] = numerator
            else:
                numerator = 0
                denominator = 0
                for i in range(k_nearest_neighbours):
                    if(k_nearest_neighbours == 1):
                        distances = neighbour_data[0][dist_lt_thresh]
                        indices = neighbour_data[1][dist_lt_thresh]
                    else:
                        distances = neighbour_data[0][dist_lt_thresh, i]
                        indices = neighbour_data[1][dist_lt_thresh, i]

                    inverse_distance = 1.0/(distances+1.0e-100)
                    values = quantity_xyValue[indices, 2]
                    numerator += values*inverse_distance
                    denominator += inverse_distance

                quantity_output[dist_lt_thresh] = numerator / denominator

        return quantity_output

    # Return the quantity function
    return quant_NN_fun

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

def composite_quantity_setting_function(poly_fun_pairs,
                                        domain,
                                        clip_range=None,
                                        nan_treatment='exception',
                                        nan_interpolation_region_polygon=None,
                                        default_k_nearest_neighbours=1,
                                        default_raster_interpolation='pixel',
                                        verbose=True):
    """ Make a 'composite function' to set quantities -- applies different
        functions inside different polygon regions.

        poly_fun_pairs = [ [p0, f0], [p1, f1], ...]

        Where:

          fi is a function,
             or a constant,
             or a '.txt' or '.csv' file with comma separated xyz data
                and an optional header row which contains letters,
             or the name of a gdal-compatible rasterFile
                (not ending in .txt or .csv),
             or a numpy array with 3 columns

          pi is a polygon (anuga polygon format),
            or a polygon filename (shapefile or a csv format that
                                    anuga.read_polygon will read),
            or None ( equivalent to a polygon with zero area),
            or 'All' (equivalent to a polygon covering everything)
            or 'Extent' in the case that fi is a rasterFile name
                (equivalent to a polygon with the same extent as the raster)

        IMPORTANT: When polygons overlap, the first elements of the list are
                   given priority. The approach is:
                   First f0 is applied to all points in p0, and we record
                     that these points have been 'set'
                   Next f1 is applied to all points in p1 which have not
                     been 'set', and then we record those points as being 'set'
                   Next f2 is applied to all points in p2 which have not
                     been 'set', and then we record those points as being 'set'
                   ... etc

        INPUT:
          @param poly_fun_pairs = [ [p0, f0], [p1, f1], ...]

              where fi(x,y) is a function returning quantity values at points,
                or any of the special cases below

              SPECIAL fi CASES:
              fi = a constant in which case points in the polygon are
                   set to that value,
              fi = a .txt or .csv file name containing x, y, z data,
                     with comma separators and an optional header row
                     containing letters (nearest neighbour interpolation is used)
              fi = a string rasterFile name (not ending in .txt or .csv)
                    which can be passed to quantityRasterFun to make a function
              fi = a numpy array with 3 columns (x,y,Value) in which case
                   nearest-neighbour interpolation is used on the points

              pi are polygons where we want to use fi inside
              (anuga polygon format) or any of the special cases below
              SPECIAL pi CASES:
              If pi is a filename ending in .shp or a csv format that
                anuga.read_polygon can read, we assume it contains a polygon
                we have to read
              If any pi = 'All', then we assume that ALL unset points are set
                 using the function. This CAN ONLY happen in the last [fi,pi]
                 pair where pi is not None (since fi will be applied to
                 all remaining points -- so anything else is probably an
                 input mistake)
              If any pi = None, then that [fi,pi] pair is skipped
              If pi = 'Extent' and fi is the name of a raster file, then the
                extent of the raster file is used to define the polygon

          @param domain = ANUGA domain object

          @param clip_range = List with the same length as poly_fun_pairs,
                 of the form:
                 [ [min0, max0], [min1, max1], ...]
                 After f0 is applied in p0, its values will be 'clipped' to the
                 range
                    [min0, max0]
                 , and similarly for the other fi

          @param nan_treatment = 'exception' or 'fall_through' -- string determining
                what to do if F(x,y) is nan. The default 'exception' raises an exception.
                The value 'fall_through' allows the function to try lower-priority
                poly,fun pairs (in sequence) to set the value.

          @param nan_interpolation_region_polygon = None, or 'All', or a list
                of csv or shp filenames containing polygons, or a list of
                anuga polygon objects.

                If it is not None, then all x,y points which evaluate to nan
                on their **first preference** dataset are recorded, and as a
                final step, the values at these x,y points
                **which are inside the nan_interpolation_region_polygon**
                are interpolated from the other x,y,F(x,y) values.

                Nearest neighbour interpolation is used, with
                k_nearest_neighbours taken from default_k_nearest_neighbours.

                Note that if nan_treatment = 'exception', then nan's will cause
                exceptions earlier on in this routine, so you will need
                nan_treatment = 'fall_through' to use this option.

                Example of why you might want this:
                    Say you have 2 elevation datasets (one defining the
                    topography above MSL, and the other defining the topography
                    below MSL). There might be small nan gaps between them,
                    which you would like to fill with interpolation. That
                    can be done with this option, by including the nan regions
                    in one of the elevation-dataset-polygons pi.

          @param default_k_nearest_neighbours = integer >=1 . The value of
                k_nearest_neighbours passed to
                make_nearestNeighbour_quantity_function when a 'special_case'
                value of fi is passed in (either a point array or a .txt or
                .csv point file), or when nan_interpolation_region_polygon is
                not None

          @param default_raster_interpolation = 'pixel' or 'bilinear'. The value of
                'interpolation' passed to quantityRasterFun if a raster filename
                is passed as one of the fi.

          @param verbose TRUE/FALSE Print more information

        OUTPUT: A function F(x,y) which can be used e.g. to set the quantity
                domain.set_quantity('elevation', F)

    """
    import os
    import numpy
    from anuga.geometry.polygon import inside_polygon

    # Check that clip_range has the right form
    if clip_range is not None:
        if len(clip_range) != len(poly_fun_pairs):
            msg = ' clip_range must be the same ' +\
                  'length as poly_fun_pairs, or None'
            raise ValueError(msg)
        # Check that min < = max
        for i in range(len(clip_range)):
            if clip_range[i][0] > clip_range[i][1]:
                raise Exception('clip_range minima must be less than maxima')

    def F(x, y):
        """This is the function returned by composite_quantity_setting_function
           It can be passed to set_quantity
        """
        isSet = numpy.zeros(len(x))  # 0/1 - record if each point has been set
        quantityVal = x*0 + numpy.nan  # Function return value

        # Record points which evaluated to nan on their first preference
        # dataset.
        was_ever_nan = (x*0).astype(int)

        lpf = len(poly_fun_pairs)
        if(lpf <= 0):
            raise Exception('Must have at least 1 fun-poly-pair')

        # Make an array of 'transformed' spatial coordinates, for checking
        # polygon inclusion
        xll = domain.geo_reference.xllcorner
        yll = domain.geo_reference.yllcorner
        xy_array_trans = numpy.vstack([x+xll, y+yll]).transpose()

        # Check that none of the pi polygons [except perhaps the last] is 'All'
        for i in range(lpf-1):
            if(poly_fun_pairs[i][0] == 'All'):
                # This is only ok if all the othe poly_fun_pairs are None
                remaining_poly_fun_pairs_are_None = \
                    [poly_fun_pairs[j][0] is None for j in range(i+1, lpf)]
                if(not all(remaining_poly_fun_pairs_are_None)):
                    raise Exception('Can only have the last polygon = All')

        # Main Loop
        # Apply the fi inside the pi
        for i in range(lpf):
            fi = poly_fun_pairs[i][1]  # The function
            pi = poly_fun_pairs[i][0]  # The polygon

            # Quick exit
            if(pi is None):
                continue

            ###################################################################
            # Get indices fInds of points in polygon pi which are not already
            # set
            ###################################################################
            if(pi == 'All'):
                # Get all unset points
                fInside = (1 - isSet)
                fInds = (fInside == 1).nonzero()[0]

            else:
                if pi == 'Extent':
                    # Here fi MUST be a gdal-compatible raster
                    if not isinstance(fi, str):
                        msg = ' pi = "Extent" can only be used when fi is a' +\
                              ' raster file name'
                        raise Exception(msg)

                    if not os.path.exists(fi):
                        msg = 'fi ' + str(fi) + ' is supposed to be a ' +\
                              ' raster filename, but it could not be found'
                        raise Exception(msg)

                    # Then we get the extent from the raster itself
                    pi_path = su.getRasterExtent(fi, asPolygon=True)

                    if verbose:
                        print('Extracting extent from raster: ', fi)
                        print('Extent: ', pi_path)

                elif (type(pi) == str) and os.path.isfile(pi):
                    # pi is a file
                    pi_path = su.read_polygon(pi)

                else:
                    # pi is the actual polygon data
                    pi_path = pi

                # Get the insides of unset points inside pi_path
                notSet = (isSet == 0.).nonzero()[0]
                fInds = inside_polygon(xy_array_trans[notSet, :], pi_path)
                fInds = notSet[fInds]

            if len(fInds) == 0:
                # No points found, move on
                continue

            ###################################################################
            # Evaluate fi at the points inside pi
            ###################################################################

            # We use various tricks to infer whether fi is a function,
            # a constant, a file (raster or csv), or an array
            if hasattr(fi, '__call__'):
                # fi is a function or a callable object
                quantityVal[fInds] = fi(x[fInds], y[fInds])

            elif isinstance(fi, (int, int, float)):
                # fi is a numerical constant
                quantityVal[fInds] = fi*1.0

            elif type(fi) is str and os.path.exists(fi):
                # fi is a file which is assumed to be
                # a gdal-compatible raster OR an x,y,z elevation file
                if os.path.splitext(fi)[1] in ['.txt', '.csv']:
                    fi_array = su.read_csv_optional_header(fi)
                    # Check the results
                    if fi_array.shape[1] != 3:
                        print('Treated input file ' + fi +
                              ' as xyz array with an optional header')
                        msg = 'Array should have 3 columns -- x,y,value'
                        raise Exception(msg)

                    newfi = make_nearestNeighbour_quantity_function(
                        fi_array, domain,
                        k_nearest_neighbours=default_k_nearest_neighbours)
                    quantityVal[fInds] = newfi(x[fInds], y[fInds])

                else:
                    # Treating input file as a raster
                    newfi = quantityRasterFun(domain, fi,
                                              interpolation=default_raster_interpolation)
                    quantityVal[fInds] = newfi(x[fInds], y[fInds])

            elif type(fi) is numpy.ndarray:
                if fi.shape[1] != 3:
                    msg = 'Array should have 3 columns -- x,y,value'
                    raise Exception(msg)
                newfi = make_nearestNeighbour_quantity_function(fi, domain,
                                                                k_nearest_neighbours=default_k_nearest_neighbours)
                quantityVal[fInds] = newfi(x[fInds], y[fInds])

            else:
                print('ERROR: with function from ' + fi)
                msg = 'Cannot make function from type ' + str(type(fi))
                raise Exception(msg)

            ###################################################################
            # Check for nan values
            ###################################################################
            #nan_flag = (quantityVal[fInds] != quantityVal[fInds])
            nan_flag = 1*numpy.isnan(quantityVal[fInds])
            nan_inds = nan_flag.nonzero()[0]
            was_ever_nan[fInds[nan_inds]] = 1

            if len(nan_inds) > 0:
                if nan_treatment == 'exception':
                    msg = 'nan values generated by the poly_fun_pair at '\
                          'index ' + str(i) + ' '\
                          'in composite_quantity_setting_function. ' + \
                          'To allow these values to be set by later ' + \
                          'poly_fun pairs, pass the argument ' + \
                          'nan_treatment="fall_through" ' + \
                          'to composite_quantity_setting_function'
                    raise Exception(msg)

                elif nan_treatment == 'fall_through':
                    msg = 'WARNING: nan values generated by the ' + \
                          'poly_fun_pair at index ' + str(i) + ' '\
                          'in composite_quantity_setting_function. ' + \
                          'They will be passed to later poly_fun_pairs'
                    if verbose:
                        print(msg)
                    not_nan_inds = (1-nan_flag).nonzero()[0]

                    if len(not_nan_inds) > 0:
                        fInds = fInds[not_nan_inds]
                    else:
                        # All values are nan
                        msg = '( Actually all the values were nan - ' + \
                              'Are you sure they should be? Possible error?)'
                        if verbose:
                            print(msg)
                        continue

                else:
                    msg = 'Found nan values in ' + \
                          'composite_quantity_setting_function but ' + \
                          'nan_treatment is not a recognized value'
                    raise Exception(msg)

            # Record that the points have been set
            isSet[fInds] = 1

            # Enforce clip_range
            if clip_range is not None:
                lower_bound = clip_range[i][0]
                upper_bound = clip_range[i][1]
                quantityVal[fInds] = numpy.maximum(
                    quantityVal[fInds], lower_bound)
                quantityVal[fInds] = numpy.minimum(
                    quantityVal[fInds], upper_bound)

        # End of loop

        # Find points which were nan on their first preference dataset + are
        # inside nan_interpolation_region_polygon. Then reinterpolate their
        # values from the other x,y, quantityVal points.
        if (nan_interpolation_region_polygon is not None) &\
           (was_ever_nan.sum() > 0):
            if nan_interpolation_region_polygon == 'All':
                points_to_reinterpolate = was_ever_nan.nonzero()[0]
            else:
                # nan_interpolation_region_polygon contains information on 1 or
                # more polygons
                # Inside those polygons, we need to re-interpolate points which
                # first evaluted to na
                possible_points_to_reint = was_ever_nan.nonzero()[0]
                points_to_reinterpolate = numpy.array([]).astype(int)

                for i in range(len(nan_interpolation_region_polygon)):
                    nan_pi = nan_interpolation_region_polygon[i]

                    # Ensure nan_pi = list of x,y points making a polygon
                    if(type(nan_pi) == str):
                        nan_pi = su.read_polygon(nan_pi)

                    points_in_nan_pi = inside_polygon(
                        xy_array_trans[possible_points_to_reint, :],
                        nan_pi)

                    if len(points_in_nan_pi) > 0:
                        points_to_reinterpolate = numpy.hstack(
                            [points_to_reinterpolate,
                             possible_points_to_reint[points_in_nan_pi]])

            if verbose:
                print('Re-interpolating ', len(points_to_reinterpolate),
                      ' points which were nan under their',
                      ' first-preference and are inside the',
                      ' nan_interpolation_region_polygon')

            if len(points_to_reinterpolate) > 0:
                msg = 'WARNING: nan interpolation is being applied. This ',\
                      'should be done in serial prior to distributing the ',\
                      'domain, as there is no parallel communication ',\
                      'implemented yet [so parallel results might depend on ',\
                      'the number of processes]'
                if verbose:
                    print(msg)

            # Find the interpolation points = points not needing reinterpolation
            ip = x*0 + 1
            ip[points_to_reinterpolate] = 0
            number_of_ip = ip.sum()
            ip = ip.nonzero()[0]

            # Check that none of the ip points has an nan value
            nan_ip = (quantityVal[ip] != quantityVal[ip]).nonzero()[0]

            if len(nan_ip) > 0:
                print('There are ', len(nan_ip), ' points outside the ',
                      'nan_interpolation_region_polygon have nan values.')
                print('The user should ensure this does not happen.')
                print('The points have the following coordinates:')
                print(xy_array_trans[ip[nan_ip], :])
                msg = "There are nan points outside of " +\
                      "nan_interpolation_region_polygon, even after all " +\
                      "fall-through's"
                raise Exception(msg)

            if(number_of_ip < default_k_nearest_neighbours):
                raise Exception('Too few non-nan points to interpolate from')

            # Make function for re-interpolation. Note this requires
            # x,y,z in georeferenced coordinates, whereas x,y are ANUGA
            # coordinates
            reinterp_F = make_nearestNeighbour_quantity_function(
                numpy.vstack([xy_array_trans[ip, 0], xy_array_trans[ip, 1],
                              quantityVal[ip]]).transpose(),
                domain,
                k_nearest_neighbours=default_k_nearest_neighbours)

            # re-interpolate
            quantityVal[points_to_reinterpolate] = reinterp_F(
                x[points_to_reinterpolate], y[points_to_reinterpolate])

            isSet[points_to_reinterpolate] = 1

        # Check there are no remaining nan values
        if(min(isSet) != 1):
            print('Some points remain as nan, which is not allowed')
            unset_inds = (isSet != 1).nonzero()[0]
            lui = min(5, len(unset_inds))
            print('There are ', len(unset_inds), ' such points')
            print('Here are a few:')
            for i in range(lui):
                print(x[unset_inds[i]] + xll, y[unset_inds[i]] + yll)
            raise Exception('It seems the input data needs to be fixed')

        return quantityVal
        # END OF FUNCTION F(x,y)

    return F

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

def quantityRasterFun(domain, rasterFile, interpolation='pixel'):
    """
    Make a function whick takes x,y in ANUGA coordinates, and returns the values
    on a raster rasterFile

    This can be used to set a quantity, and takes care of the manual conversion
    from ANUGA coordinates to spatial coordinates.

    INPUTS: @param domain = ANUGA domain
            @param rasterFile = Filename of the raster to extract point values
                    from
            @param interpolation = 'pixel' (in which case the point value is
                    set from the pixel it is on) or 'bilinear' in which case
                    the point value is set from bilinear interpolation of
                    pixels.

    OUTPUT: Function which takes x,y in ANUGA coordinates, and outputs their
            corresponding raster values
    """
    import scipy
    # import numpy as NearestNDInterpolator  # FIXME (Ole): What?
    import numpy as np

    from anuga.utilities.spatialInputUtil import rasterValuesAtPoints

    def QFun(x, y):
        xll = domain.geo_reference.xllcorner
        yll = domain.geo_reference.yllcorner
        inDat = np.vstack([x+xll, y+yll]).transpose()
        return rasterValuesAtPoints(xy=inDat, rasterFile=rasterFile,
                                    interpolation=interpolation)

    return QFun

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

def quantity_from_Pt_Pol_Data_and_Raster(Pt_Pol_Data, quantity_raster, domain):
    """
        Function to make a function F(x,y) which returns the corresponding
        values on quantity_raster, except if x,y is inside the polygon associated with
        any element of Pt_Pol_Data, in which case a Pt_Pol_-specific nearest neighbour
        interpolator is used.

        This has been superceeded by composite_quantity_setting_function

        INPUT:
            @param Pt_Pol_Data = a list with [ [ Polygon_0, Pt_XYZ_0],
                                               [ Polygon_1, Pt_XYZ_1],
                                               ...
                                             ]
                    Here Polygon_i is a polygon in ANUGA format,
                    and Pt_XYZ_i is a 3 column array of x,y,Value points
            @param quantity_raster = A GDAL-compatible quantity raster
            @param domain = ANUGA domain
    """

    # Function to set quantity from raster
    qFun1 = quantityRasterFun(domain, rasterFile=quantity_raster)

    # List of function/polygon pairs defining the Pt_Pol_ quantity data
    qFunChanList = []
    for i in range(len(Pt_Pol_Data)):
        qFunChanList.append([
            Pt_Pol_Data[i][0],
            make_nearestNeighbour_quantity_function(Pt_Pol_Data[i][1], domain)
        ])

    #
    qFun = composite_quantity_setting_function(
        qFunChanList+[['All', qFun1]], domain)

    return qFun