# Copyright (c) FULIUCANSHENG.
# Licensed under the MIT License.

import logging
import math
import time
import numpy as np
import torch
import torchvision
from torch import nn, Tensor
from typing import List, Dict, Tuple

from detectron2.config import configurable
from detectron2.layers import ShapeSpec, Conv2d
from detectron2.modeling.meta_arch.build import META_ARCH_REGISTRY
from detectron2.modeling.backbone import Backbone
from detectron2.structures import ImageList, Instances, Boxes
from detectron2.utils.events import get_event_storage
from detectron2.data.detection_utils import convert_image_to_rgb
from detectron2.modeling.postprocessing import detector_postprocess
from detectron2.modeling import build_backbone

class YoloV5Head(nn.Module):
    @configurable
    def __init__(
        self,
        *,
        input_shape: List[ShapeSpec],
        nc,
        anchors,
    ):

        super().__init__()
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        assert self.nl == len(input_shape)
        self.na = len(anchors[0]) // 2  # number of anchors
        self.grid = [torch.zeros(1)] * self.nl  # init grid
        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
        self.register_buffer("anchors", a)  # shape(nl,na,2)
        self.register_buffer("anchor_grid", a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
        ch = [x.channels for x in input_shape]
        self.m = nn.ModuleList(Conv2d(x, self.no * self.na, 1) for x in ch)

    @classmethod
    def from_config(cls, cfg, input_shape: List[ShapeSpec]):
        nc = cfg.MODEL.YOLOV5.NUM_CLASSES
        anchors = cfg.MODEL.YOLOV5.ANCHORS
        return {
            "input_shape": input_shape,
            "nc": nc,
            "anchors": anchors,
        }

    def forward(self, x: List[Tensor]):
        """
        Arguments:
            features (list[Tensor]): FPN feature map tensors in high to low resolution.
                Each tensor in the list correspond to different feature levels.
        Returns:
            x (list[Tensor]): #nl tensors,
                                each having shape [N, na, Hi, Wi, nc + 5]
            z (Tensor) : [N, nl*na*(sum of grid sizes) , no] indictaing
                    1. Box position z[..., 0:2]
                    2. Box width and height z[..., 2:4]
                    3. Objectness z[..., 5]
                    4. Class probabilities z[..., 6:]
        """
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
        return x

    @staticmethod
    def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()

    def _initialize_biases(self, cf=None):  # initialize biases into Detect(), cf is class frequency
        # https://arxiv.org/abs/1708.02002 section 3.3
        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
        for mi, s in zip(self.m, self.stride):  # from
            b = mi.bias.view(self.na, -1)  # conv.bias(255) to (3,85)
            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
            b.data[:, 5:] += math.log(0.6 / (self.nc - 0.99)) if cf is None else torch.log(cf / cf.sum())  # cls
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

def xywh2xyxy(x):
    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

def non_max_suppression(
    prediction,
    conf_thres=0.25,
    iou_thres=0.45,
    classes=None,
    agnostic=False,
    multi_label=False,
    labels=(),
    max_det=300,
):
    """Runs Non-Maximum Suppression (NMS) on inference results
        conf_thresh - 0.1 (for yolov4)
        iou_thresh - 0.6 (for yolov4)
        multi_label - not in yolov4
        merge = False - in yolov4 not in yolov3
        Labesl = () not in yolov4
    Returns:
         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
    """
    # if yolov3 or yolov5:
    nc = prediction.shape[2] - 5  # number of classes
    # else
    nc = prediction[0].shape[1] - 5  # Number of classes
    xc = prediction[..., 4] > conf_thres  # candidates

    # Checks
    assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
    assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"

    # Settings
    # (pixels) minimum and maximum box width and height
    _, max_wh = 2, 4096
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS

    t = time.time()
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        x = x[xc[xi]]  # confidence

        # Cat apriori labels if autolabelling - Not used in the YOLOV4
        if labels and len(labels[xi]):
            l_ = labels[xi]
            v = torch.zeros((len(l_), nc + 5), device=x.device)
            v[:, :4] = l_[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l_)), l_[:, 0].long() + 5] = 1.0  # cls
            x = torch.cat((x, v), 0)
        #################################################################
        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        # Detections matrix nx6 (xyxy, conf, cls)
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

        # Filter by class
        if classes is not None:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

        # Apply finite constraint
        # if not torch.isfinite(x).all():
        #     x = x[torch.isfinite(x).all(1)]

        # Check shape
        n = x.shape[0]  # number of boxes
        if not n:  # no boxes
            continue
        # #### Not in Yolov4 ######################
        elif n > max_nms:  # excess boxes
            # sort by confidence
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]
        ###############################################
        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        # boxes (offset by class), scores
        boxes, scores = x[:, :4] + c, x[:, 4]
        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
            weights = iou * scores[None]  # box weights
            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
            if redundant:
                i = i[iou.sum(1) > 1]  # require redundancy

        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            print(f"WARNING: NMS time limit {time_limit}s exceeded")
            break  # time limit exceeded

    return output

def bbox_iou(
    box1,
    box2,
    x1y1x2y2=True,
    GIoU=False,
    DIoU=False,
    CIoU=False,
    EIoU=False,
    ECIoU=False,
    eps=1e-7,
):
    # eps default value used in yolov4 is 1e-9
    # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
    box2 = box2.T

    # Get the coordinates of bounding boxes
    if x1y1x2y2:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
    else:  # transform from xywh to xyxy
        b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

    # Intersection area
    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * (
        torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)
    ).clamp(0)

    # Union Area
    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
    union = w1 * h1 + w2 * h2 - inter + eps

    iou = inter / union
    if GIoU or DIoU or CIoU or EIoU or ECIoU:
        # convex (smallest enclosing box) width
        cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)
        ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
        if CIoU or DIoU or EIoU or ECIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = (
                (b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2
            ) / 4  # center distance squared
            if DIoU:
                return iou - rho2 / c2  # DIoU
            elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU

            # ################### Function from Yolov4 ###########################################
            elif EIoU:  # Efficient IoU https://arxiv.org/abs/2101.08158
                rho3 = (w1 - w2) ** 2
                c3 = cw ** 2 + eps
                rho4 = (h1 - h2) ** 2
                c4 = ch ** 2 + eps
                return iou - rho2 / c2 - rho3 / c3 - rho4 / c4  # EIoU
            elif ECIoU:
                v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
                with torch.no_grad():
                    alpha = v / ((1 + eps) - iou + v)
                rho3 = (w1 - w2) ** 2
                c3 = cw ** 2 + eps
                rho4 = (h1 - h2) ** 2
                c4 = ch ** 2 + eps
                return iou - v * alpha - rho2 / c2 - rho3 / c3 - rho4 / c4  # ECIoU
            ############################################################################################
        else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
            c_area = cw * ch + eps  # convex area
            return iou - (c_area - union) / c_area  # GIoU
    else:
        return iou  # IoU

def smooth_BCE(
    eps=0.1,
):  # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
    # return positive, negative label smoothing BCE targets
    return 1.0 - 0.5 * eps, 0.5 * eps

class FocalLoss(nn.Module):
    # Wraps focal loss around existing loss_fcn(), i.e. criteria =
    # FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super().__init__()
        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = "none"  # required to apply FL to each element

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        # p_t = torch.exp(-loss)
        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power
        # for gradient stability

        # TF implementation
        # https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
        pred_prob = torch.sigmoid(pred)  # prob from logits
        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= alpha_factor * modulating_factor

        if self.reduction == "mean":
            return loss.mean()
        elif self.reduction == "sum":
            return loss.sum()
        else:  # 'none'
            return loss

class ComputeLoss(object):
    # Compute losses

    @configurable
    def __init__(
        self,
        *,
        focal_loss_gamma,
        box_loss_gain,
        cls_loss_gain,
        cls_positive_weight,
        obj_loss_gain,
        obj_positive_weight,
        label_smoothing=0.0,
        gr,
        na,
        nc,
        nl,
        anchors,
        anchor_t,
        autobalance=False,
    ):
        super().__init__()
        self.sort_obj_iou = False
        self.na = na
        self.nc = nc
        self.nl = nl
        self.anchors = anchors
        self.box_loss_gain = box_loss_gain
        self.cls_loss_gain = cls_loss_gain
        self.obj_loss_gain = obj_loss_gain
        self.anchor_t = anchor_t

        # Define criteria
        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([cls_positive_weight]))
        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([obj_positive_weight]))

        # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
        # positive, negative BCE targets
        self.cp, self.cn = smooth_BCE(eps=label_smoothing)

        # Focal loss
        if focal_loss_gamma > 0:
            BCEcls = FocalLoss(BCEcls, focal_loss_gamma)
            BCEobj = FocalLoss(BCEobj, focal_loss_gamma)

        # Detect() module
        self.balance = {3: [4.0, 1.0, 0.4]}.get(nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7
        self.ssi = 0  # stride 16 index
        self.BCEcls, self.BCEobj, self.gr, self.autobalance = (
            BCEcls,
            BCEobj,
            gr,
            autobalance,
        )

    @classmethod
    def from_config(cls, cfg, head):
        return {
            "focal_loss_gamma": cfg.MODEL.YOLOV5.FOCAL_LOSS_GAMMA,
            "box_loss_gain": cfg.MODEL.YOLOV5.BOX_LOSS_GAIN,
            "cls_loss_gain": cfg.MODEL.YOLOV5.CLS_LOSS_GAIN,
            "cls_positive_weight": cfg.MODEL.YOLOV5.CLS_POSITIVE_WEIGHT,
            "obj_loss_gain": cfg.MODEL.YOLOV5.OBJ_LOSS_GAIN,
            "obj_positive_weight": cfg.MODEL.YOLOV5.OBJ_POSITIVE_WEIGHT,
            "label_smoothing": cfg.MODEL.YOLOV5.LABEL_SMOOTHING,
            "gr": 1.0,
            "na": head.na,
            "nc": head.nc,
            "nl": head.nl,
            "anchors": head.anchors,
            "anchor_t": cfg.MODEL.YOLOV5.ANCHOR_T,
            "autobalance": False,
        }

    def _initialize_ssi(self, stride):
        if self.autobalance:
            self.ssi = list(stride).index(16)

    def __call__(self, p, instances):  # predictions, targets, model is ignored
        device = instances[0].gt_boxes.device
        self.to(device)
        lcls, lbox, lobj = (
            torch.zeros(1, device=device),
            torch.zeros(1, device=device),
            torch.zeros(1, device=device),
        )
        tcls, tbox, indices, anchors = self.build_targets(p, instances)  # targets

        # Losses
        for i, pi in enumerate(p):  # layer index, layer predictions
            b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
            tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj

            n = b.shape[0]  # number of targets
            if n:
                # prediction subset corresponding to targets
                ps = pi[b, a, gj, gi]

                # Regression
                pxy = ps[:, :2].sigmoid() * 2.0 - 0.5
                pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
                pbox = torch.cat((pxy, pwh), 1)  # predicted box
                # iou(prediction, target)
                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)
                lbox += (1.0 - iou).mean()  # iou loss

                # Objectness
                score_iou = iou.detach().clamp(0).type(tobj.dtype)
                if self.sort_obj_iou:
                    sort_id = torch.argsort(score_iou)
                    b, a, gj, gi, score_iou = (
                        b[sort_id],
                        a[sort_id],
                        gj[sort_id],
                        gi[sort_id],
                        score_iou[sort_id],
                    )
                tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * score_iou  # iou ratio

                # Classification
                if self.nc > 1:  # cls loss (only if multiple classes)
                    t = torch.full_like(ps[:, 5:], self.cn, device=device)  # targets
                    t[range(n), tcls[i]] = self.cp
                    lcls += self.BCEcls(ps[:, 5:], t)  # BCE

                # Append targets to text file
                # with open('targets.txt', 'a') as file:
                #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]

            obji = self.BCEobj(pi[..., 4], tobj)
            lobj += obji * self.balance[i]  # obj loss
            if self.autobalance:
                self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()

        if self.autobalance:
            self.balance = [x / self.balance[self.ssi] for x in self.balance]
        lbox *= self.box_loss_gain
        lobj *= self.obj_loss_gain
        lcls *= self.cls_loss_gain
        # bs = tobj.shape[0]  # batch size

        # loss = lbox + lobj + lcls
        # return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
        return {
            "loss_box": lbox,
            "loss_obj": lobj,
            "loss_cls": lcls,
        }

    def build_targets(self, p, gt_instances):
        """
        Args:
            p (list[Tensors]): A list of #feature level predictions
            gt_instances (list[Instances]): a list of N `Instances`s. The i-th
                `Instances` contains the ground-truth per-instance annotations
                for the i-th input image.
        """
        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
        targets = []
        for i, gt_per_image in enumerate(gt_instances):
            # Convert the boxes to target format of shape [sum(nL per image), 6]
            # where each target entry is [img_index, class, x, y, w, h],
            # x, y, w, h - relative and x, y are centers
            if len(gt_per_image) > 0:
                boxes = gt_per_image.gt_boxes.tensor.clone()
                h, w = gt_per_image.image_size
                boxes[:, 0:2] = (boxes[:, 0:2] + boxes[:, 2:4]) / 2
                boxes[:, 2:4] = (boxes[:, 2:4] - boxes[:, 0:2]) * 2
                boxes[:, ::2] /= float(w)
                boxes[:, 1::2] /= float(h)
                classes = torch.unsqueeze(gt_per_image.gt_classes.clone(), dim=1)
                t = torch.cat([torch.ones_like(classes) * i, classes, boxes], dim=1)
                targets.append(t)
        targets = torch.cat(targets, 0)

        na, nt = self.na, targets.shape[0]  # number of anchors, targets
        tcls, tbox, indices, anch = [], [], [], []
        # normalized to gridspace gain
        gain = torch.ones(7, device=targets.device)
        ai = (
            torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)
        )  # same as .repeat_interleave(nt)
        # append anchor indices
        targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)

        g = 0.5  # bias
        off = (
            torch.tensor(
                [
                    [0, 0],
                    [1, 0],
                    [0, 1],
                    [-1, 0],
                    [0, -1],  # j,k,l,m
                    # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                ],
                device=targets.device,
            ).float()
            * g
        )  # offsets

        for i in range(self.nl):
            anchors = self.anchors[i]
            gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain

            # Match targets to anchors
            t = targets * gain
            if nt:
                # Matches
                r = t[:, :, 4:6] / anchors[:, None]  # wh ratio
                j = torch.max(r, 1.0 / r).max(2)[0] < self.anchor_t  # compare
                # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  #
                # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                t = t[j]  # filter

                # Offsets
                gxy = t[:, 2:4]  # grid xy
                gxi = gain[[2, 3]] - gxy  # inverse
                j, k = ((gxy % 1.0 < g) & (gxy > 1.0)).T
                l, m = ((gxi % 1.0 < g) & (gxi > 1.0)).T
                j = torch.stack((torch.ones_like(j), j, k, l, m))
                t = t.repeat((5, 1, 1))[j]
                offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
            else:
                t = targets[0]
                offsets = 0

            # Define
            b, c = t[:, :2].long().T  # image, class
            gxy = t[:, 2:4]  # grid xy
            gwh = t[:, 4:6]  # grid wh
            gij = (gxy - offsets).long()
            gi, gj = gij.T  # grid xy indices

            # Append
            a = t[:, 6].long()  # anchor indices
            # image, anchor, grid indices
            indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))
            tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
            anch.append(anchors[a])  # anchors
            tcls.append(c)  # class

        return tcls, tbox, indices, anch

    def to(self, device):
        self.anchors = self.anchors.to(device)
        self.BCEcls.pos_weight = self.BCEcls.pos_weight.to(device)
        self.BCEobj.pos_weight = self.BCEobj.pos_weight.to(device)

@META_ARCH_REGISTRY.register()
class YoloV5(nn.Module):
    """
    Implement YoloV5
    """

    @configurable
    def __init__(
        self,
        *,
        backbone: Backbone,
        head: nn.Module,
        loss,
        num_classes,
        conf_thres,
        iou_thres,
        pixel_mean,
        pixel_std,
        vis_period=0,
        input_format="BGR",
    ):
        super().__init__()

        self.backbone = backbone
        self.head = head

        self.num_classes = num_classes
        self.single_cls = num_classes == 1
        # Inference Parameters
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        # Vis parameters
        self.vis_period = vis_period
        self.input_format = input_format

        self.register_buffer("pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), False)
        self.register_buffer("pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), False)

        """
        In Detectron1, loss is normalized by number of foreground samples in the batch.
        When batch size is 1 per GPU, #foreground has a large variance and
        using it lead to lower performance. Here we maintain an EMA of #foreground to
        stabilize the normalizer.
        """
        self.loss = loss
        # self.loss_normalizer = 100  # initialize with any reasonable #fg that's not too small
        # self.loss_normalizer_momentum = 0.9
        self.init_stride()
        self.apply(self._init_weights)

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.BatchNorm2d):
            module.eps = 1e-3
            module.momentum = 0.03

    @classmethod
    def from_config(cls, cfg):
        backbone = build_backbone(cfg)
        backbone_shape = backbone.output_shape()
        feature_shapes = list(backbone_shape.values())
        head = YoloV5Head(cfg, feature_shapes)
        loss = ComputeLoss(cfg, head)
        return {
            "backbone": backbone,
            "head": head,
            "loss": loss,
            "num_classes": head.nc,
            "conf_thres": cfg.MODEL.YOLOV5.CONF_THRESH,
            "iou_thres": cfg.MODEL.YOLOV5.IOU_THRES,
            "pixel_mean": cfg.MODEL.PIXEL_MEAN,
            "pixel_std": cfg.MODEL.PIXEL_STD,
            "vis_period": cfg.VIS_PERIOD,
            "input_format": cfg.INPUT.FORMAT,
        }

    @property
    def device(self):
        return self.pixel_mean.device

    def visualize_training(self, batched_inputs, results):
        """
        A function used to visualize ground truth images and final network predictions.
        It shows ground truth bounding boxes on the original image and up to 20
        predicted object bounding boxes on the original image.
        Args:
            batched_inputs (list): a list that contains input to the model.
            results (List[Instances]): a list of #images elements.
        """
        from detectron2.utils.visualizer import Visualizer

        assert len(batched_inputs) == len(results), "Cannot visualize inputs and results of different sizes"
        storage = get_event_storage()
        max_boxes = 20

        image_index = 0  # only visualize a single image
        img = batched_inputs[image_index]["image"]
        img = convert_image_to_rgb(img.permute(1, 2, 0), self.input_format)
        v_gt = Visualizer(img, None)
        v_gt = v_gt.overlay_instances(boxes=batched_inputs[image_index]["instances"].gt_boxes)
        anno_img = v_gt.get_image()
        processed_results = detector_postprocess(results[image_index], img.shape[0], img.shape[1])
        predicted_boxes = processed_results.pred_boxes.tensor.detach().cpu().numpy()

        v_pred = Visualizer(img, None)
        v_pred = v_pred.overlay_instances(boxes=predicted_boxes[0:max_boxes])
        prop_img = v_pred.get_image()
        vis_img = np.vstack((anno_img, prop_img))
        vis_img = vis_img.transpose(2, 0, 1)
        vis_name = f"Top: GT bounding boxes; Bottom: {max_boxes} Highest Scoring Results"
        storage.put_image(vis_name, vis_img)

    def init_stride(self):
        s = 256  # 2x min stride
        dummy_input = torch.zeros(1, len(self.pixel_mean), s, s)
        features = self.backbone(dummy_input)
        features = list(features.values())
        pred = self.head(features)
        self.head.stride = torch.tensor([s / x.shape[-2] for x in pred])  # forward
        self.head.anchors /= self.head.stride.view(-1, 1, 1)
        self.stride = self.head.stride
        self.head._initialize_biases()  # only run once
        self.loss._initialize_ssi(self.stride)

    def forward(self, batched_inputs: Tuple[Dict[str, Tensor]]):
        """
        Args:
            batched_inputs: a list, batched outputs of :class:`DatasetMapper` .
                Each item in the list contains the inputs for one image.
                For now, each item in the list is a dict that contains:
                * image: Tensor, image in (C, H, W) format.
                * instances: Instances
                Other information that's included in the original dicts, such as:
                * "height", "width" (int): the output resolution of the model, used in inference.
                  See :meth:`postprocess` for details.
        Returns:
            In training, dict[str, Tensor]: mapping from a named loss to a tensor storing the
            loss. Used during training only. In inference, the standard output format, described
            in :doc:`/tutorials/models`.
        """
        images = self.preprocess_image(batched_inputs)
        features = self.backbone(images.tensor)
        features = list(features.values())

        pred = self.head(features)

        if self.training:
            assert not torch.jit.is_scripting(), "Not supported"
            assert "instances" in batched_inputs[0], "Instance annotations are missing in training!"
            gt_instances = [x["instances"].to(self.device) for x in batched_inputs]

            losses = self.loss(pred, gt_instances)

            if self.vis_period > 0:
                storage = get_event_storage()
                if storage.iter % self.vis_period == 0:
                    results = self.inference(pred, images.image_sizes)
                    self.visualize_training(batched_inputs, results)

            return losses
        else:
            results = self.inference(pred, images.image_sizes)
            if torch.jit.is_scripting():
                return results
            processed_results = []
            for results_per_image, input_per_image, image_size in zip(results, batched_inputs, images.image_sizes):
                height = input_per_image.get("height", image_size[0])
                width = input_per_image.get("width", image_size[1])
                r = detector_postprocess(results_per_image, height, width)
                processed_results.append({"instances": r})
            return processed_results

    def inference(self, x, image_sizes):
        """
        Returns:
        z (Tensor) : [N, nl*na*(sum of grid sizes) , no] indictaing
                    1. Box position z[..., 0:2]
                    2. Box width and height z[..., 2:4]
                    3. Objectness z[..., 5]
                    4. Class probabilities z[..., 6:]
        """
        z = []
        for i in range(self.head.nl):
            # x(bs,na,ny,nx,no)
            bs, _, ny, nx, _ = x[i].shape
            if self.head.grid[i].shape[2:4] != x[i].shape[2:4]:
                self.head.grid[i] = self.head._make_grid(nx, ny).to(x[i].device)

            y = x[i].sigmoid()
            # if self.head.inplace:
            y[..., 0:2] = (y[..., 0:2] * 2.0 - 0.5 + self.head.grid[i]) * self.head.stride[i]  # xy
            y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.head.anchor_grid[i]  # wh
            # else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
            #     xy = (y[..., 0:2] * 2. - 0.5 + self.head.grid[i]) * self.head.stride[i]  # xy
            #     wh = (y[..., 2:4] * 2) ** 2 * self.head.anchor_grid[i].view(1, self.head.na, 1, 1, 2)  # wh
            #     y = torch.cat((xy, wh, y[..., 4:]), -1)
            z.append(y.view(bs, -1, self.head.no))
        return self.process_inference(torch.cat(z, 1), image_sizes)

    def process_inference(self, out, image_sizes):
        out = non_max_suppression(
            out,
            self.conf_thres,
            self.iou_thres,
            multi_label=True,
            agnostic=self.single_cls,
        )
        assert len(out) == len(image_sizes)
        results_all: List[Instances] = []
        # Statistics per image
        for si, (pred, img_size) in enumerate(zip(out, image_sizes)):

            if len(pred) == 0:
                result = Instances(img_size)
                result.pred_boxes = Boxes(torch.tensor([]))
                result.scores = torch.tensor([])
                result.pred_classes = torch.tensor([])
            else:
                # Predictions
                if self.single_cls:
                    pred[:, 5] = 0
                predn = pred.clone()
                # Predn shape [ndets, 6] of format [xyxy, conf, cls] relative to the input image size
                result = Instances(img_size)
                result.pred_boxes = Boxes(predn[:, :4])  # TODO: Check if resizing needed
                result.scores = predn[:, 4]
                result.pred_classes = predn[:, 5].int()  # TODO: Check the classes
            results_all.append(result)
        return results_all

    def preprocess_image(self, batched_inputs: Tuple[Dict[str, Tensor]]):
        """
        Normalize, pad and batch the input images.
        """
        images = [x["image"].to(self.device) for x in batched_inputs]
        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
        images = ImageList.from_tensors(images, self.backbone.size_divisibility)
        return images


"""API for joint-density-of-states calculation."""
# Copyright (C) 2019 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import numpy as np
from phonopy.structure.cells import Primitive, Supercell
from phonopy.structure.symmetry import Symmetry
from phonopy.units import VaspToTHz

from phono3py.file_IO import write_joint_dos
from phono3py.phonon3.imag_self_energy import (
    get_freq_points_batches,
    get_frequency_points,
)
from phono3py.phonon3.joint_dos import JointDos
from phono3py.phonon.grid import BZGrid

class Phono3pyJointDos:
    """Class to calculate joint-density-of-states."""

    def __init__(
        self,
        supercell: Supercell,
        primitive: Primitive,
        fc2,
        mesh=None,
        nac_params=None,
        nac_q_direction=None,
        sigmas=None,
        cutoff_frequency=1e-4,
        frequency_step=None,
        num_frequency_points=None,
        num_points_in_batch=None,
        temperatures=None,
        frequency_factor_to_THz=VaspToTHz,
        frequency_scale_factor=None,
        use_grg=False,
        SNF_coordinates="reciprocal",
        is_mesh_symmetry=True,
        is_symmetry=True,
        store_dense_gp_map=False,
        symprec=1e-5,
        output_filename=None,
        log_level=0,
    ):
        """Init method."""
        self._primitive = primitive
        self._supercell = supercell
        self._fc2 = fc2
        self._temperatures = temperatures
        self._nac_params = nac_params
        self._nac_q_direction = nac_q_direction
        if sigmas is None:
            self._sigmas = [None]
        else:
            self._sigmas = sigmas
        self._cutoff_frequency = cutoff_frequency
        self._frequency_factor_to_THz = frequency_factor_to_THz
        self._frequency_scale_factor = frequency_scale_factor
        self._is_mesh_symmetry = is_mesh_symmetry
        self._is_symmetry = is_symmetry
        self._store_dense_gp_map = store_dense_gp_map
        self._use_grg = use_grg
        self._SNF_coordinates = SNF_coordinates
        self._symprec = symprec
        self._filename = output_filename
        self._log_level = log_level

        self._bz_grid = None
        self._joint_dos = None
        self._num_frequency_points_in_batch = num_points_in_batch
        self._frequency_step = frequency_step
        self._num_frequency_points = num_frequency_points

        self._primitive_symmetry = Symmetry(
            self._primitive, self._symprec, self._is_symmetry
        )

        if mesh is not None:
            self.mesh_numbers = mesh
            self.initialize(mesh)

    @property
    def grid(self):
        """Return BZGrid class instance."""
        return self._bz_grid

    @property
    def nac_params(self):
        """Setter and getter of parameters for non-analytical term correction."""
        return self._nac_params

    @property
    def num_frequency_points_in_batch(self):
        """Getter and setter of num_frequency_points_in_batch.

        Number of sampling frequency points per batch.
        Larger value gives better concurrency in tetrahedron method,
        but requires more memory.

        """
        return self._num_frequency_points_in_batch

    @num_frequency_points_in_batch.setter
    def num_frequency_points_in_batch(self, nelems_in_batch):
        self._num_frequency_points_in_batch = nelems_in_batch

    @property
    def mesh_numbers(self):
        """Setter and getter of sampling mesh numbers in reciprocal space."""
        if self._bz_grid is None:
            return None
        else:
            return self._bz_grid.D_diag

    @mesh_numbers.setter
    def mesh_numbers(self, mesh_numbers):
        self._bz_grid = BZGrid(
            mesh_numbers,
            lattice=self._primitive.cell,
            symmetry_dataset=self._primitive_symmetry.dataset,
            is_time_reversal=self._is_symmetry,
            use_grg=self._use_grg,
            force_SNF=False,
            SNF_coordinates=self._SNF_coordinates,
            store_dense_gp_map=self._store_dense_gp_map,
        )

    def initialize(self, mesh_numbers):
        """Initialize JointDos."""
        self._jdos = JointDos(
            self._primitive,
            self._supercell,
            self._bz_grid,
            self._fc2,
            nac_params=self._nac_params,
            cutoff_frequency=self._cutoff_frequency,
            frequency_factor_to_THz=self._frequency_factor_to_THz,
            frequency_scale_factor=self._frequency_scale_factor,
            is_mesh_symmetry=self._is_mesh_symmetry,
            store_dense_gp_map=self._store_dense_gp_map,
            symprec=self._symprec,
            filename=self._filename,
            log_level=self._log_level,
        )
        if self._log_level:
            print("Generating grid system ... ", end="", flush=True)

        self.mesh_numbers = mesh_numbers

        if self._log_level:
            if self._bz_grid.grid_matrix is None:
                print("[ %d %d %d ]" % tuple(self._bz_grid.D_diag))
            else:
                print("")
                print(
                    "Generalized regular grid: [ %d %d %d ]"
                    % tuple(self._bz_grid.D_diag)
                )
                print("Grid generation matrix:")
                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[0]))
                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[1]))
                print("  [ %d %d %d ]" % tuple(self._bz_grid.grid_matrix[2]))

    def run(self, grid_points, write_jdos=False):
        """Calculate joint-density-of-states."""
        if self._log_level:
            print(
                "--------------------------------- Joint DOS "
                "---------------------------------"
            )
            print("Running harmonic phonon calculations...", flush=True)

        self._jdos.run_phonon_solver()
        frequencies, _, _ = self._jdos.get_phonons()
        self._jdos.run_phonon_solver_at_gamma()
        max_phonon_freq = np.max(frequencies)
        self._jdos.run_phonon_solver_at_gamma(is_nac=True)

        self._frequency_points = get_frequency_points(
            max_phonon_freq=max_phonon_freq,
            sigmas=self._sigmas,
            frequency_points=None,
            frequency_step=self._frequency_step,
            num_frequency_points=self._num_frequency_points,
        )
        batches = get_freq_points_batches(
            len(self._frequency_points), nelems=self._num_frequency_points_in_batch
        )
        if self._temperatures is None:
            temperatures = [None]
        else:
            temperatures = self._temperatures
        self._joint_dos = np.zeros(
            (
                len(self._sigmas),
                len(temperatures),
                len(self._frequency_points),
                2,
            ),
            dtype="double",
            order="C",
        )

        for i, gp in enumerate(grid_points):
            if (self._bz_grid.addresses[gp] == 0).all():
                self._jdos.nac_q_direction = self._nac_q_direction
            else:
                self._jdos.nac_q_direction = None
            self._jdos.set_grid_point(gp)

            if self._log_level:
                weights = self._jdos.get_triplets_at_q()[1]
                print(
                    "======================= "
                    "Grid point %d (%d/%d) "
                    "=======================" % (gp, i + 1, len(grid_points))
                )
                adrs = self._jdos.bz_grid.addresses[gp]
                q = np.dot(adrs, self._bz_grid.QDinv.T)
                print("q-point: (%5.2f %5.2f %5.2f)" % tuple(q))
                print("Number of triplets: %d" % len(weights))
                print("Frequency")
                for f in self._jdos.get_phonons()[0][gp]:
                    print("%8.3f" % f)

            if not self._sigmas:
                raise RuntimeError("sigma or tetrahedron method has to be set.")

            for i_s, sigma in enumerate(self._sigmas):
                self._jdos.sigma = sigma
                if self._log_level:
                    if sigma is None:
                        print("Tetrahedron method is used.")
                    else:
                        print("Smearing method with sigma=%s is used." % sigma)
                    print(
                        f"Calculations at {len(self._frequency_points)} "
                        f"frequency points are devided into {len(batches)} batches."
                    )
                for i_t, temperature in enumerate(temperatures):
                    self._jdos.temperature = temperature

                    for ib, freq_indices in enumerate(batches):
                        if self._log_level:
                            print(
                                f"{ib + 1}/{len(batches)}: {freq_indices + 1}",
                                flush=True,
                            )
                        self._jdos.frequency_points = self._frequency_points[
                            freq_indices
                        ]
                        self._jdos.run()
                        self._joint_dos[i_s, i_t, freq_indices] = self._jdos.joint_dos

                    if write_jdos:
                        filename = self._write(gp, i_sigma=i_s)
                        if self._log_level:
                            print('JDOS is written into "%s".' % filename)

    @property
    def dynamical_matrix(self):
        """Return DynamicalMatrix class instance."""
        return self._jdos.dynamical_matrix

    @property
    def frequency_points(self):
        """Return frequency points."""
        return self._frequency_points

    @property
    def joint_dos(self):
        """Return calculated joint-density-of-states."""
        return self._joint_dos

    def _write(self, gp, i_sigma=0):
        return write_joint_dos(
            gp,
            self._bz_grid.D_diag,
            self._frequency_points,
            self._joint_dos[i_sigma],
            sigma=self._sigmas[i_sigma],
            temperatures=self._temperatures,
            filename=self._filename,
            is_mesh_symmetry=self._is_mesh_symmetry,
        )


from abc import ABC
from typing import Union

import numpy as np
from qcore import geo
from qcore.nhm import NHMFault

from srf_generation.source_parameter_generation.uncertainties.mag_scaling import (
    MagnitudeScalingRelations,
    get_area,
    get_length,
    get_width,
    mag2mom,
    round_subfault_size,
    lw_to_mw_scaling_relation,
)

LEONARD_SEISMOGENIC_DEPTH_DIFFERENCE = 3
NHM_SEISMOGENIC_DEPTH = 12

def fault_factory(fault_type: int):
    return [Type1, Type2, Type3, Type4][fault_type - 1]

class Fault(ABC):
    subfault_spacing = 0.1
    type = 0
    _mag = None
    _moment = None
    name = None
    _latitude = None
    _longitude = None
    _strike = None
    _rake = None
    _dip = None
    _depth = None
    _width = None
    _length = None
    _shypo = None
    _dhypo = None

    @property
    def pid(self):
        return self.name

    @pid.setter
    def pid(self, value):
        self.name = value

    @property
    def mom(self):
        if self._moment is None:
            self._moment = mag2mom(self._mag)
        return self._moment

    @property
    def latitude(self):
        return self._latitude

    @latitude.setter
    def latitude(self, value):
        self._latitude = value

    @property
    def longitude(self):
        return self._longitude

    @longitude.setter
    def longitude(self, value):
        self._longitude = value

    @property
    def magnitude(self):
        return self._mag

    @magnitude.setter
    def magnitude(self, mag):
        if mag > 11:
            raise ValueError(
                f"Given mag {mag} is greater than theoretically possible 11"
            )
        self._mag = mag

    @property
    def length(self):
        return self._length

    @property
    def width(self):
        return self._width

    def to_dict(self):
        return {
            "type": self.type,
            "magnitude": self._mag,
            "moment": self.mom,
            "name": self.name,
            "longitude": self.longitude,
            "latitude": self.latitude,
            "strike": self._strike,
            "rake": self._rake,
            "dip": self._dip,
            "depth": self._depth,
        }

    @property
    def rake(self):
        return self._rake

    @rake.setter
    def rake(self, value):
        value = ((value + 180) % 360) - 180
        self._set_rake(value)

    def _set_rake(self, value):
        self._rake = value

    @property
    def dip(self):
        return self._dip

    @dip.setter
    def dip(self, value):
        if 90 < value < 180:
            self.strike = self.strike + 180
            self.rake = -self.rake
            value = 180 - value
        elif value < 0 or value > 180:
            # The fault is now above ground
            raise ValueError(f"Invalid dip value: {value}")

        self._set_dip(value)

    def _set_dip(self, value):
        self._dip = value

class SinglePlaneFault(Fault):
    def __init__(self, name, magnitude, strike, rake, dip):
        self.name = name
        self._mag = magnitude
        self._strike = strike
        self._rake = rake
        self._dip = dip

    _length: float = None
    _width: float = None
    _dtop: float = None
    _dbottom: float = None

    @property
    def strike(self):
        return self._strike

    @strike.setter
    def strike(self, value):
        self._strike = value % 360

    @property
    def dhypo(self):
        return self.width / 2

    @property
    def dbottom(self):
        return self._dbottom

    @property
    def dtop(self):
        return self._dtop

    @property
    def hypocentre_lonlat(self):
        return self._longitude, self._latitude

class MultiPlaneFault(Fault):
    pass

class PointSourceFault(SinglePlaneFault):

    type = 1

    # vs = 3.2
    # rh0 = 2.44
    # risetime = 0.5
    # stype = "cos"
    # inittime = 0.0
    vs = None
    rh0 = None
    risetime = None
    stype = None
    inittime = None

    def __init__(self, name, latitude, longitude, magnitude, strike, rake, dip, depth):
        self._depth = depth
        self._latitude = latitude
        self._longitude = longitude
        super().__init__(name, magnitude, strike, rake, dip)

    @property
    def dbottom(self):
        return self._depth

    def to_dict(self):
        base_dict = super().to_dict()
        base_dict.update({"type": 1})
        if self.vs is not None:
            base_dict["vs"] = self.vs
        if self.rh0 is not None:
            base_dict["rho"] = self.rh0
        if self.risetime is not None:
            base_dict["risetime"] = self.risetime
        if self.stype is not None:
            base_dict["stype"] = self.stype
        if self.inittime is not None:
            base_dict["inittime"] = self.inittime
        return base_dict

Type1 = PointSourceFault

class FiniteFault(SinglePlaneFault):

    dlen = 0.1
    dwid = 0.1
    shypo = 0.00
    depths = None

    def to_dict(self):
        base_dict: dict = super().to_dict()
        base_dict.update(
            {
                "flen": self.length,
                "dlen": self.dlen,
                "fwid": self.width,
                "dwid": self.dwid,
                "dtop": self.dtop,
                "shypo": self.shypo,
                "dhypo": self.dhypo,
            }
        )
        return base_dict

class Type2(FiniteFault):

    type = 2

    _ratio_override: Union[float, None] = None
    _magnitude_scaling_relation: MagnitudeScalingRelations = None

    def __init__(self, name, latitude, longitude, magnitude, strike, rake, dip, depth):
        self._depth = depth
        self._latitude = latitude
        self._longitude = longitude
        super().__init__(name, magnitude, strike, rake, dip)

    def to_dict(self):
        base_dict: dict = super().to_dict()
        base_dict.update(
            {
                "type": 2,
                "longitude": self._lon_hyp,
                "latitude": self._lat_hyp,
                "mwsr": self._magnitude_scaling_relation.name,
                "dbottom": self.dbottom,
            }
        )
        return base_dict

    def _set_rake(self, value):
        self._rake = value
        if self._magnitude_scaling_relation is not None:
            self._calculate_dimensions()

    def set_ratio_override(self, ratio):
        self._ratio_override = ratio

    def reset_ratio_override(self):
        self._ratio_override = None

    @property
    def magnitude_scaling_relation(self):
        return self._magnitude_scaling_relation

    @magnitude_scaling_relation.setter
    def magnitude_scaling_relation(self, value: MagnitudeScalingRelations):
        self._magnitude_scaling_relation = value
        self._calculate_dimensions()

    @property
    def length(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")

        return self._length

    @property
    def width(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")

        return self._width

    @property
    def dhypo(self):
        return self.width / 2

    @property
    def dbottom(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        return self._dbottom

    @property
    def dtop(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        return self._dtop

    @property
    def hypocentre_lonlat(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        return self._lon_hyp, self._lat_hyp

    def _calculate_dimensions(self):
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")

        area = get_area(self)
        length = get_length(self)
        width = get_width(self)

        if self._ratio_override is not None:
            r = self._ratio_override
        else:
            r = max(length / width, 1)

        self._width = np.sqrt(area / r)
        self._length = r * np.sqrt(area / r)

        self._dtop = self._depth - np.sin(np.radians(self._dip)) * self.width / 2
        shift = min(self._dtop, 0)
        if shift != 0:
            self._dtop = 0
        self._dbottom = (
            self._depth + np.sin(np.radians(self._dip)) * self.width / 2 + shift
        )
        if (
            self.magnitude_scaling_relation == MagnitudeScalingRelations.LEONARD2014
            and self._dbottom > NHM_SEISMOGENIC_DEPTH
        ):
            self._dbottom += LEONARD_SEISMOGENIC_DEPTH_DIFFERENCE

        self.ny = int(round(self.width / self.dwid))
        self.nx = int(round(self.length / self.dlen))
        self.dlen = self.length / float(self.nx)
        self.dwid = self.width / float(self.ny)

        # self._calculate_finite_fault_properties()
        # self._calculate_corners()
        self._get_hypocentre(0, -self.dhypo)

    def _get_hypocentre(self, shypo, dhypo):
        """
        Same logic as corners, for a single point.
        """
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        ONE_DEG_LAT = np.radians(6371.0072)
        azimuth = np.radians(self.strike - 90)
        rot_matrix = np.array(
            [[np.cos(azimuth), np.sin(azimuth)], [-np.sin(azimuth), np.cos(azimuth)]]
        )

        y_pos_surf_proj_hyp = -dhypo * np.cos(np.radians(self._dip))
        hypocentre_points = np.dot(rot_matrix, [[shypo], [y_pos_surf_proj_hyp]])

        self._lat_hyp = self.latitude + hypocentre_points[1][0] / ONE_DEG_LAT
        self._lon_hyp = self.longitude + (
            hypocentre_points[0][0] / ONE_DEG_LAT
        ) / np.cos(np.radians(self.latitude))

    def _calculate_corners(self):
        # use cartesian coordinate system to define the along strike and downdip
        # locations taking the center of the fault plane as (x,y)=(0,0)
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        ONE_DEG_LAT = np.radians(6371.0072)

        x_pos = np.arange(self.dlen / 2.0, self.length, self.dlen) - self.length / 2.0
        y_pos = (
            np.arange(self.dwid / 2.0, self.width, self.dwid)[::-1] - self.width / 2.0
        )

        # now use a coordinate transformation to go from fault plane to North and
        # East cartesian plane (again with (0,0) ==(0,0)  )
        y_pos_surf_proj = y_pos * np.cos(np.radians(self._dip))
        azimuth = np.radians(self.strike - 90)
        rot_matrix = np.array(
            [[np.cos(azimuth), np.sin(azimuth)], [-np.sin(azimuth), np.cos(azimuth)]]
        )
        east_loc_relative, north_loc_relative = np.dot(
            rot_matrix, [np.tile(x_pos, self.ny), y_pos_surf_proj.repeat(self.nx)]
        ).reshape((2, self.nx, self.ny))
        # now use a coordinate transformation to go from the North East cartesian
        # plane to the spherical earth WGS84 coordinate system
        self.lats = self.latitude + north_loc_relative / ONE_DEG_LAT
        self.lons = self.longitude + (east_loc_relative / ONE_DEG_LAT) * 1 / np.cos(
            np.radians(self.latitude)
        )

    def _calculate_finite_fault_properties(self):
        """
        Purpose: To create a finite fault geometry based on centroid moment tensor
        solution information and magnitude scaling relationships.
        The use of such finite fault geometry is for first order computation of
        source-to-site distances.

        revisions
        v2 - reoriented fault plane to be consistent with along strike direction
        v3 - added 'argout_emod3d' output and associated commands to output details
             which are input in the finite fault model generation for EMOD3D graves methodology

        assumptions:
        1) The centroid moment tensor represents the centroid of the fault plane,
        located half way along strike and downdip
        2) The fault area is computed from the median of a Moment scaling
        relationship and thus is assumed to be a 'typical' stress drop event for
        the Mw-relationship used.
        Should only be used for small events where the downdip width is smaller
        than the seismogenic depth (i.e. so the assumption of a square fault plane
        is reasonable); depth is reset to zero if above ground values determined
        3) srike in deg measured clockwise from N; rake in deg measured
        anti-clockwise from the strike direction (and in the range
        -180<lambda<180; dip in deg measured downward from horizontal

        Input variables:
        Lat    - the latitude of the focal mechanism (-ve below equator)
        Lon    - the lon of the focal mech (-180<lon<180)
        Depth  - the depth of the focal mech (in km)
        Mw     - the moment magnitude from the Mw tensor soln
        strike - the strike of the focal mech (only 1)
        rake   - rake (not used, but carried forward)
        dip    - the dip angle of the fault plane

        output variables:
        lat       - the latitude of the subfault
        lon       - the lon of the subfault
        depth     - depth of the subfault
        lonAsList - as a 1D array
        """
        if self.magnitude_scaling_relation is None:
            raise ValueError("No magnitude scaling relation given")
        # rounded subfault spacing

        # use cartesian coordinate system to define the along strike and downdip
        # locations taking the center of the fault plane as (x,y)=(0,0)
        y_pos: np.ndarray = (
            np.arange(self.dwid / 2.0, self.width, self.dwid)[::-1] - self.width / 2.0
        )

        depth_loc_relative = (
            (-y_pos * np.sin(np.radians(self._dip)))
            .repeat(self.nx)
            .reshape((self.nx, self.ny))
        )
        shift = np.min(self.dtop + depth_loc_relative, 0)
        self.depths = self.dtop + depth_loc_relative - shift
        if np.any(0 > self.depths):
            raise ValueError(
                "Some points are above the ground. This represents a logic problem and should be referred to the "
                "developers."
            )

class Type3(FiniteFault):
    def __init__(
        self,
        name,
        magnitude,
        lon1,
        lat1,
        lon2,
        lat2,
        dip,
        rake,
        dtop,
        dbottom,
        dip_dir,
        tectonic_type,
    ):
        self._trace = ((lon1, lat1), (lon2, lat2))
        self._dtop = dtop
        self._dbottom = dbottom
        self._dip_dir = dip_dir

        self._length = round_subfault_size(
            geo.ll_dist(lon1, lat1, lon2, lat2), magnitude
        )

        if tectonic_type == "SUBDUCTION_INTERFACE":
            self.mwsr = MagnitudeScalingRelations.SKARLATOUDIS2016

        else:
            self.mwsr = MagnitudeScalingRelations.LEONARD2014
            self._dbottom += 3

        raw_fwid = (self._dbottom - dtop) / np.sin(np.radians(dip))
        self._width = round_subfault_size(raw_fwid, magnitude)

        strike = geo.ll_bearing(lon1, lat1, lon2, lat2)

        if (self._dip_dir - strike + 360) % 360 > 180:
            # Reverse the order of the points so the dipdir is to the right of the strike
            self._trace = self._trace[::-1]
            strike = geo.ll_bearing(lon2, lat2, lon1, lat1)

        super().__init__(name, magnitude, strike, rake, dip)

    def to_dict(self):
        base_dict = {
            "clon": geo.ll_mid(*self._trace[0], *self._trace[1])[0],
            "clat": geo.ll_mid(*self._trace[0], *self._trace[1])[1],
            "type": self.type,
            "magnitude": self._mag,
            "moment": self.mom,
            "name": self.name,
            "strike": self._strike,
            "rake": self._rake,
            "dip": self._dip,
            "dtop": self._dtop,
            "dbottom": self._dbottom,
            "length": self._length,
            "width": self._width,
            "dip_dir": self._dip_dir,
        }
        return base_dict

class Type4(MultiPlaneFault):

    type = 4

    def __init__(self, nhm_data: NHMFault):
        self.name = nhm_data.name
        self._dbottom = nhm_data.dbottom
        self.fault_type = nhm_data.fault_type
        self.tectonic_type = nhm_data.tectonic_type
        self._dip = nhm_data.dip
        self._rake = nhm_data.rake
        self._dtop = nhm_data.dtop
        self._slip_rate = nhm_data.slip_rate
        self._dip_dir = nhm_data.dip_dir

        self._n_planes = len(nhm_data.trace) - 1

        if nhm_data.tectonic_type == "SUBDUCTION_INTERFACE":
            self.mwsr = MagnitudeScalingRelations.SKARLATOUDIS2016

        else:
            self.mwsr = MagnitudeScalingRelations.LEONARD2014
            self._dbottom += 3

        dummy_plane = Type3(
            nhm_data.name,
            nhm_data.mw,
            *nhm_data.trace[0],
            *nhm_data.trace[1],
            nhm_data.dip,
            nhm_data.rake,
            nhm_data.dtop,
            nhm_data.dbottom,
            nhm_data.dip_dir,
            nhm_data.tectonic_type,
        )

        length = sum(
            [
                geo.ll_dist(*nhm_data.trace[i], *nhm_data.trace[i + 1])
                for i in range(self._n_planes)
            ]
        )

        self._mag = lw_to_mw_scaling_relation(
            length, dummy_plane.width, self.mwsr, nhm_data.rake
        )

        if dummy_plane.strike != geo.ll_bearing(*nhm_data.trace[0], *nhm_data.trace[1]):
            nhm_data.trace = nhm_data.trace[::-1]

        self._planes = []
        for i in range(self._n_planes):
            self._planes.append(
                Type3(
                    nhm_data.name,
                    self._mag,
                    *nhm_data.trace[i],
                    *nhm_data.trace[i + 1],
                    nhm_data.dip,
                    nhm_data.rake,
                    nhm_data.dtop,
                    nhm_data.dbottom,
                    nhm_data.dip_dir,
                    nhm_data.tectonic_type,
                )
            )

    @property
    def length(self):
        return sum([f.length for f in self._planes])

    @property
    def width(self):
        return self._planes[0].width

    @property
    def dhypo(self):
        return self._dhypo

    @dhypo.setter
    def dhypo(self, dhypo):
        if dhypo > self.width:
            raise ValueError(
                f"Cannot place hypocentre outside fault plane. dhpyo: {dhypo}, fault width: {self.width}"
            )
        self._dhypo = dhypo
        for sub_plane in self._planes:
            sub_plane._dhypo = dhypo

    @property
    def shypo(self):
        return self._shypo

    @shypo.setter
    def shypo(self, shypo):
        if shypo > self.length:
            raise ValueError(
                f"Cannot place hypocentre outside fault plane. shpyo: {shypo}, fault length: {self.length}"
            )
        self._shypo = shypo
        for sub_plane in self._planes:
            sub_plane._shypo = shypo

    def to_dict(self):

        base_dict = {
            "type": self.type,
            "magnitude": self._mag,
            "moment": self.mom,
            "name": self.name,
            "fault_type": self.fault_type,
            "tect_type": self.tectonic_type,
            "rake": self._rake,
            "dip": self._dip,
            "dtop": self._dtop,
            "dbottom": self._dbottom,
            "length": self.length,
            "plane_count": self._n_planes,
            "slip_rate": self._slip_rate,
            "dip_dir": self._dip_dir,
            "shypo": self.shypo,
            "dhypo": self.dhypo,
        }
        for i, sub_fault in enumerate(self._planes):
            for key, item in sub_fault.to_dict().items():
                base_dict[f"{key}_subfault_{i}"] = item
        return base_dict

    @property
    def trace(self):
        points = [x._trace[0] for x in self._planes]
        points.append(self._planes[-1]._trace[1])
        return np.asarray(points)

    @trace.setter
    def trace(self, values):
        for i in range(self._n_planes):
            sub_plane = self._planes[i]
            sub_plane._trace = (tuple(values[i]), tuple(values[i + 1]))

    @property
    def dtop(self):
        return self._dtop

    @dtop.setter
    def dtop(self, value):
        if value < 0:
            raise ValueError("dtop must be below ground level")
        self._dtop = value

    @property
    def dbottom(self):
        return self._dbottom

    @dbottom.setter
    def dbottom(self, value):
        if value < 0:
            raise ValueError("dtop must be below ground level")
        self._dbottom = value

    @property
    def dip_dir(self):
        return self._dip_dir

    @dip_dir.setter
    def dip_dir(self, value):
        self._dip_dir = value

    def _set_dip(self, value):
        for sub_fault in self._planes:
            sub_fault.dip = value
        self._dip = value

    @property
    def strike(self):
        return np.asarray([x.strike for x in self._planes])

    @strike.setter
    def strike(self, values):
        for i in range(self._n_planes):
            self._planes[i].strike = values[i]


#!/usr/bin/python3
#
##########################################################################################################################
#                                                                                #
#                   What follows are options of the PP-STM code for calculating inner PDOS                               #
#                                                                                #
##########################################################################################################################
#
# Note : This type of simulations works for solid slabs or molecules on slabs (substrate) ; for freestanding molecule it can give you nonsences
#
# ***** System information: *****
#
ppstm_path = './PPSTM/'      # path (absolute or relative) to your PPSTM code #
#
# ***** Main informations ******
#
V_min         = -2.0         # V_min for plotting
V_max         = +2.0         # V_max for plotting #
eta           =  0.1         # Lorentzian width of states in energy scale: typically 0.1; can be in range of 0.3-0.05 eV in some cases (low amount of layers ...) even up to 1.0 eV #
dV            =  eta/5.0     # voltage/energy step , default : eta/5.0  #
sample_orbs   = 'spd'        # orbitals of the sample 'sp' (light atoms only, faster) or 'spd' (all atoms) #
dft_code      = 'AIMS'       # 'fireball'='Fireball'='FIREBALL' ; 'aims'='AIMS'='FHI-AIMS' ; 'cp2k'='CP2K' ; 'gpaw'='GPAW' #
geometry_file = 'geom-cube.in'   # E.G. 'input.xyz' , 'input.bas' , 'geometry.in'; None for GPAW #
spin          =  'both'      # None=False ; for FHI-AIMS & CP2K: None -- spin-unpolarized/spin-restricted calc. ;  'both'--  up will be plotted as + and down as - , 'up'='alpha' or 'down"='beta' (last 3 are for spin-polarizes or spin-unrestricted calculations) #
cp2k_name     = 'FePc'       # Name used in CP2K calculations or GPAW calc #
#
# *****Plotting options ******
#
red_line_atoms    = [0]     # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; [0] -- means 1st atom #
red_line_shell    = 'dz2'   # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
blue_line_atoms   = [0,1]     # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; [1] -- means 2nd atom #
blue_line_shell   = 'all'     # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
green_line_atoms  = [0]     # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; range(2,8) = [2,3,4,5,6,7]-- means atoms: 3,4,5,6,7 and 8 #
green_line_shell  = 'all'   # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
orange_line_atoms = list(range(1,57))   # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; [8,9] -- means 9th and 10th atom #
orange_line_shell = 'pz'    # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
black_line_atoms  = list(range(1,57))   # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; None -- no line #
black_line_shell  = 'all'   # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
yellow_line_atoms = None    # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; None -- no line #
yellow_line_shell = 'all'   # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
gray_line_atoms   = None    # atoms to plot: !!! DO NOT FORGET PYTHON NUMBERING starts from 0 !!!, always use list-like structure; 'all' = -1 -- all atoms; None -- no line #
gray_line_shell   = 'all'   # which chells to plot: 'all' = all sample orbitals, 's', 'p', 'd', 'px', 'py', 'pz', 'dxy', 'dyz', 'dyz', 'dz2', 'dxz', 'dx2y2' #
#
# *****Output options ******
#
PNG  = True                  # True / False -- plot "png" images (2D graph height) #
TXT = False                  # **** Not working yet *** True / False -- write ".txt" files with 1st column energy, other columns PDOS#
#
# ***** Advanced options ******
#
cut_atoms   = None           # Not really neccessary for these calculations, unless really big files and memory; None = -1 -- All atoms of the sample contributes to tunelling ; 1 -- only 1st atom of the sample contributes to the tunelling ; 57 -- first 57 atoms of the sample contributes to the tunelling ; ... #
#
# ***** More advanced options ******
#
fermi        = None          # None=0.0 -- no change to the Fermi Level ; -0.1 -- shifts the Fermi Level by 0.1 eV lower ... #
cut_min      = V_min-3*eta   # cut out all orbitals lower than  -default: V_min-3*eta bellow Fermi (should be: cut_min <= Vmin-2*eta) . taken to the Fermi Level #
cut_max      = V_max+3*eta   # cut out all orbitals higher than -default: V_max+3*eta above  Fermi (should be: cut_max >= Vmax+2*eta) . taken to the Fermi Level #
files_path   = ''            # where are files fron DFT code ; rather do not use this #
lower_atoms  = 'no-d-rescalling' # normally d-orbs are rescalled by factor of 0.2 #
#
#
##########################################################################################################################
#                                                                                #
#                                 DO NOT TOUCH LATER CODE (unless you know what you are doing)                           #
#                                                                                #
##########################################################################################################################

print("Importing libraries")

import os
import sys
sys.path.append(ppstm_path)

import numpy as np
#import pyPPSTM                   as PS
import pyPPSTM.ReadSTM           as RS
import pyPPSTM.PreSTMutils       as SU
import matplotlib
matplotlib.use('Agg') # Force matplotlib to not use any Xwindows backend. ## !!! important for working on clusters !!!!
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
#if (plot_atoms):
#    import pyPPSTM.basUtils as Bu
#    import pyPPSTM.elements as elements

print("Libraries imported")

# --- some function definition --- #

def printf(*args):
    together = ''.join(map(str, args))    # avoid the arg is not str
    #print together
    return together

# --- Initial check --- #

assert( PNG or TXT ), "No output set to be True; I'm not going to do anything if there is no output. I'm too lazy like a Gartfield. "

# --- reading of the eigen-energies, the LCAO coefficients and geometry --- #

print("Reading electronic & geometry structure files")

cell=[[0,0],[0,0]];pbc=(0,0);lower_coefs=[];

if ((dft_code == 'fireball') or(dft_code == 'Fireball') or (dft_code == 'FIREBALL')):
    eigEn, coefs, Ratin = RS.read_FIREBALL_all(name = files_path + 'phik_0001_', geom=files_path+geometry_file, lvs = cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max,cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs);

elif ((dft_code == 'gpaw') or(dft_code == 'GPAW')):
    eigEn, coefs, Ratin = RS.read_GPAW_all(    name = files_path + cp2k_name + '.gpw', fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs);

elif ((dft_code == 'aims') or(dft_code == 'AIMS') or (dft_code == 'FHI-AIMS')):
    if ((spin == None) or (spin == False)):
        name = 'KS_eigenvectors.band_1.kpt_1.out'
    elif ((spin == 'up')or(spin == 'alpha')or(spin == 'both')):
        name = 'KS_eigenvectors_up.band_1.kpt_1.out'
    elif ((spin == 'down')or(spin == 'beta')or(spin == 'dn')):
        name = 'KS_eigenvectors_dn.band_1.kpt_1.out'
    else :
        print("unknown spin, I'm going to sleep. Good Night"); exit()
    eigEn, coefs, Ratin = RS.read_AIMS_all(name = files_path + name , geom= files_path + geometry_file, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs);
    if (spin == 'both'):
        name = 'KS_eigenvectors_dn.band_1.kpt_1.out'
        eigEn2, coefs2, Ratin = RS.read_AIMS_all(name = files_path + name , geom= files_path + geometry_file, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs);
        #coefs2 *= -1;

elif ((dft_code == 'cp2k') or(dft_code == 'CP2K')):
    if ((spin == None)or(spin == False)):
        eigEn, coefs, Ratin  = RS.read_CP2K_all(name = files_path + cp2k_name , lvs=cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs);
    elif ((spin == 'up')or(spin == 'alpha')):
        eigEn, coefs, Ratin  = RS.read_CP2K_all(name = files_path + cp2k_name , lvs=cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs, spin='alpha');
    elif (spin == 'both'):
        eigEn, coefs, Ratin  = RS.read_CP2K_all(name = files_path + cp2k_name , lvs=cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs, spin='alpha');
        eigEn2, coefs2, Ratin  = RS.read_CP2K_all(name = files_path + cp2k_name , lvs=cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs, spin='beta');
        #coefs2 *= -1;
    elif ((spin == 'down')or(spin == 'beta')or(spin == 'dn')):
        eigEn, coefs, Ratin  = RS.read_CP2K_all(name = files_path + cp2k_name , lvs=cell, fermi=fermi, orbs = sample_orbs, pbc=pbc, cut_min=cut_min, cut_max=cut_max, cut_at=cut_atoms, lower_atoms=lower_atoms, lower_coefs=lower_coefs, spin='beta');
    else :
        print("unknown spin, I'm going to sleep. Good Night"); exit()

#print "DEBUG: eigEn.shape ", eigEn.shape
#print "DEBUG: coefs.shape ", coefs.shape
#print "DEBUG: Ratin.shape ", Ratin.shape

energies = np.arange(V_min,V_max,dV)

print("energies prepared, coeffecients read")

# --- the PDOS calculations --- #

plot = True if ( (red_line_atoms != None)or(blue_line_atoms != None)or(green_line_atoms != None)or(orange_line_atoms != None)or(black_line_atoms != None)or(yellow_line_atoms != None)or (gray_line_atoms != None) ) else False

print("DEBUG: plot", plot)

assert plot!=False, "No lines to plot"

latomslegend = 4

if (red_line_atoms != None) :
    atoms=red_line_atoms
    PDOS1 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms   , spherical=red_line_shell   )
    red_line_legend     = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
    print("DEBUG: red_line_legend", red_line_legend)
else:	#(red_line_atoms != None) :
    PDOS1 = None

if (blue_line_atoms != None) :
    atoms=blue_line_atoms
    PDOS2 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms  , spherical=blue_line_shell  )
    blue_line_legend    = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(blue_line_atoms != None) :
    PDOS2 = None

if (green_line_atoms != None) :
    atoms=green_line_atoms
    PDOS3 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms , spherical=green_line_shell )
    green_line_legend   = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(green_line_atoms != None) :
    PDOS3 = None

if (orange_line_atoms != None) :
    atoms=orange_line_atoms
    PDOS4 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms, spherical=orange_line_shell)
    orange_line_legend  = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(orange_line_atoms != None) :
    PDOS4 = None

if (black_line_atoms != None) :
    atoms=black_line_atoms
    PDOS5 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms , spherical=black_line_shell )
    black_line_legend   = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(black_line_atoms != None) :
    PDOS5 = None

if (yellow_line_atoms != None) :
    atoms=yellow_line_atoms
    PDOS6 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms, spherical=yellow_line_shell)
    yellow_line_legend  = ','.join(map(str, atoms)) if len(atoms) < latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(yellow_line_atoms != None) :
    PDOS6 = None

if (gray_line_atoms != None) :
    atoms=gray_line_atoms
    PDOS7 = SU.pPDOS(eigEn,coefs, energies, eta=eta, orbs= sample_orbs, atoms=atoms  , spherical=gray_line_shell  )
    gray_line_legend    = ','.join(map(str, atoms)) if len(atoms) > latomslegend else str(atoms[0])+".."+str(atoms[-1])
else:	#(gray_line_atoms != None) :
    PDOS7 = None

if spin=='both' :
    print("DEBUG: printing spin down in both spins")

    PDOS1d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=red_line_atoms   , spherical=red_line_shell   ) if (red_line_atoms != None) else None
    PDOS2d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=blue_line_atoms  , spherical=blue_line_shell  ) if (blue_line_atoms != None) else None
    PDOS3d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=green_line_atoms , spherical=green_line_shell ) if (green_line_atoms != None) else None
    PDOS4d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=orange_line_atoms, spherical=orange_line_shell) if (orange_line_atoms != None) else None
    PDOS5d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=black_line_atoms , spherical=black_line_shell ) if (black_line_atoms != None) else None
    PDOS6d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=yellow_line_atoms, spherical=yellow_line_shell) if (yellow_line_atoms != None) else None
    PDOS7d = -1*SU.pPDOS(eigEn2,coefs2, energies, eta=eta, orbs= sample_orbs, atoms=gray_line_atoms  , spherical=gray_line_shell  ) if (gray_line_atoms != None) else None

else:
    print("DEBUG: both spins are not used")
    PDOS1d = PDOS2d = PDOS3d = PDOS4d = PDOS5d = PDOS6d = PDOS7d = None

# --- plotting part here, plots all calculated signals --- #

if PDOS1 is not None:
    plt.plot(energies, PDOS1, ls='-',c='r'      ,label= 'PDOS atoms:'+red_line_legend+";"+red_line_shell+"-shell")
if PDOS2 is not None:
    plt.plot(energies, PDOS2, ls='-',c='g'      ,label= 'PDOS atoms:'+blue_line_legend+";"+blue_line_shell+"-shell")
if PDOS3 is not None:
    plt.plot(energies, PDOS3, ls='-',c='b'      ,label= 'PDOS atoms:'+green_line_legend+";"+green_line_shell+"-shell")
if PDOS4 is not None:
    plt.plot(energies, PDOS4, ls='-',c='#FFA500',label= 'PDOS atoms:'+orange_line_legend+";"+orange_line_shell+"-shell")
if PDOS5 is not None:
    plt.plot(energies, PDOS5, ls='-',c='k'      ,label= 'PDOS atoms:'+black_line_legend+";"+black_line_shell+"-shell")
if PDOS6 is not None:
    plt.plot(energies, PDOS6, ls='-',c='y'      ,label= 'PDOS atoms:'+yellow_line_legend+";"+yellow_line_shell+"-shell")
if PDOS7 is not None:
    plt.plot(energies, PDOS7, ls='-',c='g'      ,label= 'PDOS atoms:'+gray_line_legend+";"+gray_line_shell+"-shell")

# spin down if "both"
if PDOS1d is not None:
    plt.plot(energies, PDOS1d, ls='--',c='r')#,label= 'PDOS atoms:'+str(red_line_atoms)+";"+red_line_shell+"shell")
if PDOS2d is not None:
    plt.plot(energies, PDOS2d, ls='--',c='g')#,label= 'PDOS atoms:'+str(blue_line_atoms)+";"+red_line_shell+"shell")
if PDOS3d is not None:
    plt.plot(energies, PDOS3d, ls='--',c='b')#,label= 'PDOS atoms:'+str(green_line_atoms)+";"+red_line_shell+"shell")
if PDOS4d is not None:
    plt.plot(energies, PDOS4d, ls='--',c='#FFA500')#,label= 'PDOS atoms:'+str(orange_line_atoms)+";"+red_line_shell+"shell")
if PDOS5d is not None:
    plt.plot(energies, PDOS5d, ls='--',c='k')#,label= 'PDOS atoms:'+str(red_line_atoms)+";"+red_line_shell+"shell")
if PDOS6d is not None:
    plt.plot(energies, PDOS6d, ls='--',c='y')#,label= 'PDOS atoms:'+str(red_line_atoms)+";"+red_line_shell+"shell")
if PDOS7d is not None:
    plt.plot(energies, PDOS7d, ls='--',c='g')#,label= 'PDOS atoms:'+str(red_line_atoms)+";"+red_line_shell+"shell")

if (plot and PNG) :
    plt.title("PDOS eta: "+str(eta)+' eV')
    plt.xlabel("E-Efermi [eV]")
    plt.ylabel("DOS [arb.un]")
    plt.legend(loc='center left',bbox_to_anchor=(1,0.5))
    plt.savefig('PDOS_eta_'+str(eta)+'eV_orbs_'+sample_orbs+'.png', bbox_inches='tight')

# --- the end --- #

print()
print()
print("Done")
print()