module opt_orient
    !This module contains the orient option which allows for the reorientation 
    !of simulation cells. This can be used to create arbitrarily oriented dislocation or loops.
    use parameters
    use subroutines
    use elements
    use box
    implicit none


    real(kind=dp), save :: new_orient(3,3)
    real(kind=dp), dimension(6) :: orig_box_bd
    real(kind=dp), allocatable :: orig_sub_box_ori(:,:,:)
    character(len=3) :: orig_box_bc 

    public 
    contains

    subroutine orient_opt(arg_pos)

        integer, intent(inout) :: arg_pos

        integer :: i, j, k, ibasis, inod
        logical :: isortho, isrighthanded, matching
        real(kind=dp) :: inv_sub_box_ori(3,3,sub_box_num)


        !First parse the orient command
        call parse_orient(arg_pos)

        !Now rotate the basis. To do this we transform the basis to [100] [010] [001] and then 
        !transform to user specified basis.

        !Find all inverse orientation matrices for all sub_boxes
        do i = 1, sub_box_num
            call matrix_inverse(sub_box_ori(:,:,i), 3, inv_sub_box_ori(:,:,i))
        end do

        !Now transform all atoms
        do i = 1, atom_num
            r_atom(:,i) = matmul(new_orient,matmul(inv_sub_box_ori(:,:,sbox_atom(i)),r_atom(:,i)))
        end do 

        !Now transform all elements
        do i = 1, ele_num
            do inod =1, ng_node(lat_ele(i))
                do ibasis = 1, basisnum(lat_ele(i))
                    r_node(:,ibasis,inod,i) = matmul(new_orient,matmul(inv_sub_box_ori(:,:,sbox_ele(i)),r_node(:,ibasis,inod,i)))
                end do
            end do  
        end do

        !Now save the original sub_box_ori and overwrite them
        if(allocated(orig_sub_box_ori)) deallocate(orig_sub_box_ori)

        allocate(orig_sub_box_ori(3,3,sub_box_num))
        orig_sub_box_ori = sub_box_ori

        !Now overwrite the orientations
        do i = 1, sub_box_num
            sub_box_ori(:,:,i) = new_orient
        end do 

        !Save original box boundaries
        orig_box_bd = box_bd

        !Now find new box boundaries, if any orientations are the same we leave them as they are. If they are different then we have
        !to shrink wrap them

        orig_box_bc = box_bc
        do i = 1,3
            matching=.true.
            sbox_loop:do j = 1, sub_box_num
                do k = 1, 3
                    if(.not.is_equal(orig_sub_box_ori(i,k,j), new_orient(i,k))) then 
                        matching = .false.
                        exit sbox_loop
                    end if
                end do
            end do sbox_loop
            if(.not.matching) box_bc(i:i)='s'
        end do

            call def_new_box
    end subroutine orient_opt

    subroutine parse_orient(arg_pos)
        !This parses the orient option 
        integer, intent(inout) :: arg_pos

        integer :: i, arg_len
        logical :: isortho, isrighthanded
        character(len=20) :: ori_string

        !Pull out the new user orientation
        do i = 1, 3
            arg_pos = arg_pos + 1
            call get_command_argument(arg_pos, ori_string, arg_len)
            if (arg_len == 0) print *, "Missing orientation vector in -orient option"
            call parse_ori_vec(ori_string, new_orient(i,:))
        end do

        !Normalize the orientation matrix
        new_orient = matrix_normal(new_orient,3)

        !Check right hand rule and orthogonality
        call check_right_ortho(new_orient, isortho, isrighthanded)
        if (.not.isortho) then 
            stop "Directions in orient are not orthogonal"
        else if (.not.isrighthanded) then 
            stop "Directions in orient are not righthanded"
        end if

        arg_pos = arg_pos + 1

    end subroutine parse_orient

    subroutine unorient

        integer :: i, ibasis, inod
        real(kind=dp) :: inv_ori(3,3)

        !Now rotate the basis. To do this we transform the basis to [100] [010] [001] and then 
        !transform to the original sbox_ele

        !Find the inverse for the new orientation matrix
        call matrix_inverse(new_orient, 3, inv_ori)

        !Recover original sub_box_ori
        sub_box_ori = orig_sub_box_ori 

        !Now transform all atoms
        do i = 1, atom_num
            r_atom(:,i) = matmul(sub_box_ori(:,:,sbox_atom(i)),matmul(inv_ori(:,:),r_atom(:,i)))
        end do 

        !Now transform all elements
        do i = 1, ele_num
            do inod =1, ng_node(lat_ele(i))
                do ibasis = 1, basisnum(lat_ele(i))
                    r_node(:,ibasis,inod,i) = matmul(sub_box_ori(:,:,sbox_ele(i)),matmul(inv_ori,r_node(:,ibasis,inod,i)))
                end do
            end do  
        end do

        !Restore original box boundaries and box BC
        box_bd = orig_box_bd
        box_bc = orig_box_bc
    end subroutine unorient

    subroutine sbox_ori(arg_pos)
        integer, intent(inout) :: arg_pos

        integer :: i, sbox_in, arg_len
        real(kind = dp) :: new_orient(3,3)
        character(len=100) :: textholder 
        character(len=20) :: ori_string

        arg_pos = arg_pos + 1
        call get_command_argument(arg_pos, textholder,arg_len)
        if (arg_len== 0) stop 'Missing sbox in sbox_ori command'
        read(textholder,*) sbox_in

        do i = 1, 3
            arg_pos = arg_pos + 1
            call get_command_argument(arg_pos, ori_string, arg_len)
            if (arg_len == 0) print *, "Missing orientation vector in -orient option"
            call parse_ori_vec(ori_string, new_orient(i,:))
            new_orient(i,:) = new_orient(i,:) / norm2(new_orient(i,:))
        end do

        sub_box_ori(:,:,sbox_in) = new_orient
        arg_pos = arg_pos + 1

    end subroutine sbox_ori
end module opt_orient