"""Обратный расчёт девиатора напряжений по механизмам очагов землетрясений.

Линейная схема Майкла (Michael, 1984): касательное напряжение на плоскости
разрыва, вычисленное для подбираемого девиатора, должно совпадать по
направлению с наблюдаемым смещением в очаге (гипотеза Уоллеса--Ботта).
Условие совпадения линейно по компонентам тензора, поэтому девиатор
определяется методом наименьших квадратов.

Единственная внешняя библиотека -- NumPy. Случайные числа и бутстреп
не используются: при фиксированных входных данных результат единственный.

Каталог событий и контрольные значения -- Barth & Wenzel (2010).

"""
import numpy as np


def mech(strike, dip, rake):
    """Единичные векторы нормали n и смещения s по простиранию, падению, скольжению."""
    p, d, l = np.radians([strike, dip, rake])
    n = np.array([-np.sin(d)*np.sin(p),  np.sin(d)*np.cos(p), -np.cos(d)])
    s = np.array([np.cos(l)*np.cos(p) + np.sin(l)*np.cos(d)*np.sin(p),
                  np.cos(l)*np.sin(p) - np.sin(l)*np.cos(d)*np.cos(p),
                  -np.sin(l)*np.sin(d)])
    return n/np.linalg.norm(n), s/np.linalg.norm(s)


def shear(T, n):
    """Касательное напряжение на плоскости с нормалью n: tau = T n - (n . T n) n."""
    tn = T @ n
    return tn - np.dot(n, tn)*n


def beta_deg(T, n, s):
    """Невязка beta: угол между смещением s и направлением касательного напряжения."""
    tau = shear(T, n)
    return np.degrees(np.arccos(np.clip(np.dot(s, tau/np.linalg.norm(tau)), -1, 1)))


def swap(ns):
    """Узловая пара в обратном порядке: (n, s) -> (s, n)."""
    return (ns[1], ns[0])


# Ниже выписаны пять базисных тензоров; подстановка каждого в shear() даёт столбец системы МНК.
BASIS = [
    np.array([[1, 0, 0], [0, 0, 0], [0, 0, -1]]),   # компонента T11
    np.array([[0, 0, 0], [0, 1, 0], [0, 0, -1]]),   # компонента T22
    np.array([[0, 1, 0], [1, 0, 0], [0, 0, 0]]),    # компонента T12
    np.array([[0, 0, 1], [0, 0, 0], [1, 0, 0]]),    # компонента T13
    np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0]]),    # компонента T23
]


def tensor_from_components(t5):
    """Матрица 3 x 3 девиатора из пяти независимых компонент."""
    T11, T22, T12, T13, T23 = t5
    return np.array([[T11, T12, T13], [T12, T22, T23], [T13, T23, -T11 - T22]])


def solve_lstsq(data):
    """МНК-решение системы shear(T, n) = s по совокупности пар (n, s)."""
    blocks = []
    b = []
    for n, s in data:
        columns = []
        for E in BASIS:
            columns.append(shear(E, n))          # реакция базисного тензора на площадке
        blocks.append(np.column_stack(columns))  # блок 3 x 5 для одного наблюдения
        b.append(s)
    solution = np.linalg.lstsq(np.vstack(blocks), np.concatenate(b), rcond=None)
    t5 = solution[0]                             # пять независимых компонент девиатора
    T = tensor_from_components(t5)
    dots = []                                    # проверка знака: сжатие должно быть > 0
    for n, s in data:
        dots.append(np.dot(s, shear(T, n)))
    if np.mean(dots) < 0:
        T = -T
    return T


def azimuth_deg(u):
    """Азимут горизонтальной проекции оси, ° (0...180, отсчёт от севера)."""
    return np.degrees(np.arctan2(u[1], u[0])) % 180


def plunge_deg(u):
    """Угол погружения оси, ° (0...90)."""
    return np.degrees(np.arcsin(abs(u[2])))


def invert(events, max_iter=100):
    """Двухступенчатый обратный расчёт; возвращает тензор, невязки и параметры."""
    m = []
    for e in events:
        m.append(mech(e[0], e[1], e[2]))

    both = []
    for ns in m:
        both.append(ns)            # первая узловая плоскость события
        both.append(swap(ns))      # вторая (перевёрнутая пара)
    T = solve_lstsq(both)          # 1-я ступень: обе узловые плоскости

    choice_prev = None
    for _iteration in range(max_iter):   # 2-я ступень: плоскость с меньшей невязкой
        choice = []
        for ns in m:
            b_first = beta_deg(T, ns[0], ns[1])    # невязка первой плоскости
            b_second = beta_deg(T, ns[1], ns[0])   # невязка второй плоскости
            if b_first <= b_second:
                choice.append(0)
            else:
                choice.append(1)
        if choice_prev is not None and choice == choice_prev:
            break                  # выбор плоскостей стабилизировался
        choice_prev = choice
        chosen = []
        for ns, c in zip(m, choice):
            if c == 0:
                chosen.append(ns)
            else:
                chosen.append(swap(ns))
        T = solve_lstsq(chosen)

    betas = []                     # каждой паре -- меньшая из двух невязок
    for ns in m:
        betas.append(min(beta_deg(T, ns[0], ns[1]), beta_deg(T, ns[1], ns[0])))

    w, v = np.linalg.eigh(T)       # собственное разложение девиатора
    lam_a, lam_b, lam_c = w[2], w[1], w[0]
    R = (lam_a - lam_b)/(lam_a - lam_c)    # = (sigma2 - sigma3)/(sigma1 - sigma3)
    s1, s3 = v[:, 0], v[:, 2]              # оси наибольшего сжатия и растяжения
    return T, betas, dict(SH=azimuth_deg(s1), R=R, mean_beta=np.mean(betas),
                          sigma1=(azimuth_deg(s1), plunge_deg(s1)),
                          sigma3=(azimuth_deg(s3), plunge_deg(s3)))


def beta_event(T, event):
    """Невязка события относительно готового тензора (минимум по двум плоскостям)."""
    n, s = mech(event[0], event[1], event[2])
    return min(beta_deg(T, n, s), beta_deg(T, s, n))


def misfit_with_event(events, event):
    """Невязка события относительно тензора, полученного с его участием.

    Так в публикации оценены невязки исключённых выбросов: тензор строится
    по выборке, содержащей проверяемое событие (обе его плоскости), затем для
    каждого механизма оставляется плоскость с меньшей невязкой и тензор
    перевычисляется один раз. 
    """
    m = []
    for e in list(events) + [event]:
        m.append(mech(e[0], e[1], e[2]))
    both = []
    for ns in m:
        both.append(ns)
        both.append(swap(ns))
    T = solve_lstsq(both)
    chosen = []
    for ns in m:
        if beta_deg(T, ns[0], ns[1]) <= beta_deg(T, ns[1], ns[0]):
            chosen.append(ns)
        else:
            chosen.append(swap(ns))
    T = solve_lstsq(chosen)
    n, s = mech(event[0], event[1], event[2])
    return min(beta_deg(T, n, s), beta_deg(T, s, n))


# Каталог механизмов Barth & Wenzel (2010), таблица 1:
# номер: (широта, долгота, простирание, падение, скольжение)
CATALOG = {
 1:(53.5,112.2,148,47,3),   2:(47.9,130.7,33,81,-172), 3:(47.5,117.0,128,57,20),
 4:(51.6,133.4,122,32,52),  5:(49.0,129.9,200,33,131), 6:(44.7,112.6,135,44,69),
 7:(55.0,124.0,217,54,136), 8:(48.9,131.5,127,65,-5),  9:(48.9,131.2,9,71,-174),
 10:(43.4,117.5,111,45,-15),11:(44.7,115.7,40,55,-164),12:(54.5,124.2,325,42,149),
 13:(48.5,128.6,152,29,66), 14:(49.0,130.4,263,41,-1), 15:(49.4,119.4,125,54,17),
 16:(49.0,130.0,115,50,14), 17:(47.4,116.3,207,43,177),18:(44.6,117.5,225,72,162),
 19:(51.1,124.8,107,48,8),  20:(47.7,117.0,35,22,-144),21:(54.0,134.3,192,60,178),
 22:(55.4,124.2,114,43,-21),23:(43.8,119.7,313,66,4),  24:(53.9,134.4,91,51,-11),
 25:(43.4,120.2,212,65,-178),26:(53.2,128.9,97,65,5),  27:(45.4,118.3,157,27,68),
 28:(54.4,125.5,120,53,-3), 29:(48.9,131.4,300,63,-8), 30:(51.8,122.6,207,57,-173),
 31:(55.2,122.9,114,47,-17),32:(48.7,132.5,101,63,-3), 33:(48.5,131.8,191,39,127),
 34:(46.9,124.9,154,47,63), 35:(51.8,116.3,127,73,5),  36:(44.6,124.2,156,32,48),
 37:(43.5,119.6,298,69,-5), 38:(48.8,133.4,180,46,162),39:(53.6,132.4,132,56,24),
 40:(49.2,122.4,198,57,162),41:(54.1,128.0,108,64,-1),
 42:(44.7,126.9,318,48,-10),43:(47.9,130.6,104,74,-10),44:(43.8,125.1,42,64,156),
 45:(51.1,135.3,188,39,32), 46:(54.3,126.5,25,80,-179),47:(54.4,125.4,247,35,-20),
 48:(53.6,132.2,132,24,96), 49:(43.9,114.2,286,59,-17),50:(54.1,122.0,316,21,141),
 51:(54.2,128.9,110,84,0),  52:(55.2,124.8,112,27,90), 53:(48.7,121.0,130,54,6),
 54:(48.6,126.1,116,63,-12),
}


def mechanism(nr):
    """Механизм события №nr из каталога: (простирание, падение, скольжение)."""
    return CATALOG[nr][2:]        # первые два числа -- координаты очага


def mechanisms(numbers):
    """Механизмы событий с указанными номерами каталога."""
    result = []
    for nr in numbers:
        result.append(mechanism(nr))
    return result


BIN3 = [2, 5, 8, 9, 13, 16, 29, 32, 33, 38, 43]                 # Буреинский блок
BIN2 = [7, 12, 21, 22, 24, 26, 28, 31, 39, 41, 46, 48, 50, 51, 52]  # Становой район
EV2013 = (303, 75, -5)                                          # C201307221508A

if __name__ == '__main__':
    print('1. Буреинский блок, N = 11')
    print('   публикация: SH = 70.2, R = 0.40, средняя невязка = 12.2')
    T3, b3, o3 = invert(mechanisms(BIN3))
    print(f'   расчёт:     SH = {o3["SH"]:.1f}, R = {o3["R"]:.2f}, '
          f'средняя невязка = {o3["mean_beta"]:.1f}')

    print('2. Становой район, N = 15')
    print('   публикация: SH = 59.1, R = 0.47, средняя = 15.0, наибольшая < 32')
    T2, b2, o2 = invert(mechanisms(BIN2))
    print(f'   расчёт:     SH = {o2["SH"]:.1f}, R = {o2["R"]:.2f}, '
          f'средняя = {o2["mean_beta"]:.1f}, наибольшая = {max(b2):.1f}')

    print('3. Контрольные невязки; протокол публикации:')
    print('   тензор получен с участием проверяемого события')
    b47 = misfit_with_event(mechanisms(BIN2), mechanism(47))
    b45 = misfit_with_event(mechanisms(BIN2 + [4]), mechanism(45))
    print(f'   событие 47 (исключено из становой выборки): '
          f'публикация 96, расчёт {b47:.1f}')
    print(f'   событие 45 (становая выборка с переходной зоной B): '
          f'публикация 63, расчёт {b45:.1f}')
    print(f'   событие 28 (в становой выборке): '
          f'публикация 9.1, расчёт {beta_event(T2, mechanism(28)):.1f}')

    print('4. Рабочая выборка: Буреинский блок + событие 2013 г., N = 12')
    T, b, o = invert(mechanisms(BIN3) + [EV2013])
    print(f'   SH = {o["SH"]:.1f}, R = {o["R"]:.2f}, средняя = {o["mean_beta"]:.1f}, '
          f'наибольшая = {max(b):.1f}')
    print(f'   невязка события 2013 г. относительно буреинского тензора: '
          f'{beta_event(T3, EV2013):.1f}')
    T2p, b2p, o2p = invert(mechanisms(BIN2) + [EV2013])
    print(f'   относительно станового тензора: {beta_event(T2, EV2013):.1f}, '
          f'при включении в становую выборку: {b2p[-1]:.1f}')