from __future__ import absolute_import, division, print_function
import iotbx.ncs
import iotbx.ncs as ncs
from iotbx import pdb

pdb_str_1="""\
CRYST1  421.520  421.520  421.520  90.00  90.00  90.00 P 1
SCALE1      0.002372  0.000000  0.000000        0.00000
SCALE2      0.000000  0.002372  0.000000        0.00000
SCALE3      0.000000  0.000000  0.002372        0.00000
00 20.00           O
ATOM  20750  N   GLN A3048     172.895 272.881 181.246  1.00 20.00           N
ATOM  20751  CA  GLN A3048     174.038 273.205 180.405  1.00 20.00           C
ATOM  20752  C   GLN A3048     174.922 274.193 181.146  1.00 20.00           C
ATOM  20753  O   GLN A3048     174.493 275.310 181.448  1.00 20.00           O
ATOM  20754  CB  GLN A3048     173.603 273.813 179.071  1.00 20.00           C
ATOM  20755  CG  GLN A3048     172.982 272.857 178.082  1.00 20.00           C
ATOM  20756  CD  GLN A3048     172.668 273.524 176.749  1.00 20.00           C
ATOM  20757  OE1 GLN A3048     172.822 274.735 176.593  1.00 20.00           O
ATOM  20758  NE2 GLN A3048     172.245 272.728 175.776  1.00 20.00           N
ATOM  20759  N   ARG A3049     176.156 273.792 181.437  1.00 20.00           N
ATOM  20760  CA  ARG A3049     177.074 274.640 182.189  1.00 20.00           C
ATOM  20761  C   ARG A3049     178.316 274.918 181.353  1.00 20.00           C
ATOM  20762  O   ARG A3049     179.235 274.074 181.309  1.00 20.00           O
ATOM  20763  CB  ARG A3049     177.431 274.001 183.534  1.00 20.00           C
ATOM  20764  CG  ARG A3049     178.396 274.815 184.382  1.00 20.00           C
ATOM  20765  CD  ARG A3049     177.873 276.220 184.584  1.00 20.00           C
ATOM  20766  NE  ARG A3049     176.541 276.220 185.179  1.00 20.00           N
ATOM  20767  CZ  ARG A3049     175.728 277.269 185.183  1.00 20.00           C
ATOM  20768  NH1 ARG A3049     176.111 278.407 184.624  1.00 20.00           N
ATOM  20769  NH2 ARG A3049     174.529 277.182 185.738  1.00 20.00           N
ATOM  20770  N   PRO A3050     178.391 276.070 180.668  1.00 20.00           N
ATOM  20771  CA  PRO A3050     179.609 276.516 179.987  1.00 20.00           C
ATOM  20772  C   PRO A3050     180.606 277.102 180.972  1.00 20.00           C
ATOM  20773  O   PRO A3050     180.209 277.351 182.108  1.00 20.00           O
ATOM  20774  CB  PRO A3050     179.099 277.587 179.023  1.00 20.00           C
ATOM  20775  CG  PRO A3050     177.888 278.128 179.689  1.00 20.00           C
ATOM  20776  CD  PRO A3050     177.254 276.969 180.399  1.00 20.00           C
HETATM20777  C1' FMN A3994     169.244 267.550 126.473  1.00 20.00           C
HETATM20778  C10 FMN A3994     168.218 265.806 127.908  1.00 20.00           C
HETATM20779  C2  FMN A3994     168.985 263.559 127.660  1.00 20.00           C
HETATM20780  C2' FMN A3994     168.655 267.365 125.077  1.00 20.00           C
HETATM20781  C3' FMN A3994     169.586 267.961 124.044  1.00 20.00           C
HETATM20782  C4  FMN A3994     167.291 264.058 129.235  1.00 20.00           C
HETATM20783  C4' FMN A3994     169.743 269.452 124.272  1.00 20.00           C
HETATM20784  C4A FMN A3994     167.336 265.394 128.881  1.00 20.00           C
HETATM20785  C5' FMN A3994     170.460 270.146 123.129  1.00 20.00           C
HETATM20786  C5A FMN A3994     166.562 267.625 129.128  1.00 20.00           C
HETATM20787  C6  FMN A3994     165.727 268.539 129.744  1.00 20.00           C
HETATM20788  C7  FMN A3994     165.784 269.873 129.378  1.00 20.00           C
HETATM20789  C7M FMN A3994     164.880 270.860 130.041  1.00 20.00           C
HETATM20790  C8  FMN A3994     166.673 270.291 128.396  1.00 20.00           C
HETATM20791  C8M FMN A3994     166.765 271.726 127.976  1.00 20.00           C
HETATM20792  C9  FMN A3994     167.504 269.374 127.783  1.00 20.00           C
HETATM20793  C9A FMN A3994     167.449 268.040 128.146  1.00 20.00           C
HETATM20794  N1  FMN A3994     169.038 264.885 127.302  1.00 20.00           N
HETATM20795  N10 FMN A3994     168.274 267.125 127.533  1.00 20.00           N
HETATM20796  N3  FMN A3994     168.110 263.141 128.630  1.00 20.00           N
HETATM20797  N5  FMN A3994     166.508 266.302 129.495  1.00 20.00           N
HETATM20798  O1P FMN A3994     168.602 272.266 122.870  1.00 20.00           O
HETATM20799  O2  FMN A3994     169.722 262.750 127.115  1.00 20.00           O
HETATM20800  O2' FMN A3994     168.523 265.992 124.811  1.00 20.00           O
HETATM20801  O2P FMN A3994     170.055 273.770 124.179  1.00 20.00           O
HETATM20802  O3' FMN A3994     170.842 267.349 124.198  1.00 20.00           O
HETATM20803  O3P FMN A3994     170.661 273.201 121.851  1.00 20.00           O
HETATM20804  O4  FMN A3994     166.516 263.678 130.091  1.00 20.00           O
HETATM20805  O4' FMN A3994     168.489 270.061 124.431  1.00 20.00           O
HETATM20806  O5' FMN A3994     170.809 271.411 123.629  1.00 20.00           O
HETATM20807  P   FMN A3994     170.022 272.701 123.117  1.00 20.00           P
TER
ATOM  41557  N   GLN B3048     175.987 146.935 181.245  1.00 20.00           N
ATOM  41558  CA  GLN B3048     175.135 147.763 180.404  1.00 20.00           C
ATOM  41559  C   GLN B3048     173.837 148.035 181.145  1.00 20.00           C
ATOM  41560  O   GLN B3048     173.084 147.105 181.447  1.00 20.00           O
ATOM  41561  CB  GLN B3048     174.825 147.082 179.070  1.00 20.00           C
ATOM  41562  CG  GLN B3048     175.963 147.022 178.081  1.00 20.00           C
ATOM  41563  CD  GLN B3048     175.543 146.417 176.748  1.00 20.00           C
ATOM  41564  OE1 GLN B3048     174.417 145.945 176.592  1.00 20.00           O
ATOM  41565  NE2 GLN B3048     176.443 146.448 175.775  1.00 20.00           N
ATOM  41566  N   ARG B3049     173.568 149.304 181.436  1.00 20.00           N
ATOM  41567  CA  ARG B3049     172.374 149.675 182.188  1.00 20.00           C
ATOM  41568  C   ARG B3049     171.513 150.612 181.352  1.00 20.00           C
ATOM  41569  O   ARG B3049     171.784 151.830 181.308  1.00 20.00           O
ATOM  41570  CB  ARG B3049     172.750 150.304 183.533  1.00 20.00           C
ATOM  41571  CG  ARG B3049     171.562 150.733 184.381  1.00 20.00           C
ATOM  41572  CD  ARG B3049     170.607 149.578 184.583  1.00 20.00           C
ATOM  41573  NE  ARG B3049     171.273 148.424 185.178  1.00 20.00           N
ATOM  41574  CZ  ARG B3049     170.770 147.196 185.182  1.00 20.00           C
ATOM  41575  NH1 ARG B3049     169.593 146.958 184.623  1.00 20.00           N
ATOM  41576  NH2 ARG B3049     171.445 146.201 185.737  1.00 20.00           N
ATOM  41577  N   PRO B3050     170.477 150.101 180.667  1.00 20.00           N
ATOM  41578  CA  PRO B3050     169.482 150.933 179.986  1.00 20.00           C
ATOM  41579  C   PRO B3050     168.476 151.504 180.972  1.00 20.00           C
ATOM  41580  O   PRO B3050     168.459 151.036 182.108  1.00 20.00           O
ATOM  41581  CB  PRO B3050     168.809 149.956 179.023  1.00 20.00           C
ATOM  41582  CG  PRO B3050     168.946 148.637 179.689  1.00 20.00           C
ATOM  41583  CD  PRO B3050     170.267 148.667 180.398  1.00 20.00           C
HETATM41584  C1' FMN B3994     182.422 146.436 126.471  1.00 20.00           C
HETATM41585  C10 FMN B3994     184.446 146.419 127.906  1.00 20.00           C
HETATM41586  C2  FMN B3994     186.008 148.207 127.657  1.00 20.00           C
HETATM41587  C2' FMN B3994     182.877 146.019 125.075  1.00 20.00           C
HETATM41588  C3' FMN B3994     181.895 146.527 124.042  1.00 20.00           C
HETATM41589  C4  FMN B3994     186.423 146.490 129.232  1.00 20.00           C
HETATM41590  C4' FMN B3994     180.525 145.918 124.270  1.00 20.00           C
HETATM41591  C4A FMN B3994     185.243 145.861 128.879  1.00 20.00           C
HETATM41592  C5' FMN B3994     179.565 146.192 123.127  1.00 20.00           C
HETATM41593  C5A FMN B3994     183.698 144.076 129.126  1.00 20.00           C
HETATM41594  C6  FMN B3994     183.324 142.896 129.742  1.00 20.00           C
HETATM41595  C7  FMN B3994     182.140 142.279 129.376  1.00 20.00           C
HETATM41596  C7M FMN B3994     181.737 141.002 130.039  1.00 20.00           C
HETATM41597  C8  FMN B3994     181.333 142.840 128.394  1.00 20.00           C
HETATM41598  C8M FMN B3994     180.044 142.202 127.974  1.00 20.00           C
HETATM41599  C9  FMN B3994     181.712 144.018 127.781  1.00 20.00           C
HETATM41600  C9A FMN B3994     182.895 144.637 128.144  1.00 20.00           C
HETATM41601  N1  FMN B3994     184.833 147.590 127.300  1.00 20.00           N
HETATM41602  N10 FMN B3994     183.275 145.809 127.531  1.00 20.00           N
HETATM41603  N3  FMN B3994     186.808 147.658 128.627  1.00 20.00           N
HETATM41604  N5  FMN B3994     184.871 144.690 129.493  1.00 20.00           N
HETATM41605  O1P FMN B3994     178.658 143.523 122.868  1.00 20.00           O
HETATM41606  O2  FMN B3994     186.341 149.249 127.112  1.00 20.00           O
HETATM41607  O2' FMN B3994     184.132 146.591 124.809  1.00 20.00           O
HETATM41608  O2P FMN B3994     176.629 144.030 124.178  1.00 20.00           O
HETATM41609  O3' FMN B3994     181.797 147.921 124.196  1.00 20.00           O
HETATM41610  O3P FMN B3994     176.819 144.839 121.850  1.00 20.00           O
HETATM41611  O4  FMN B3994     187.140 146.009 130.088  1.00 20.00           O
HETATM41612  O4' FMN B3994     180.624 144.527 124.429  1.00 20.00           O
HETATM41613  O5' FMN B3994     178.295 145.862 123.627  1.00 20.00           O
HETATM41614  P   FMN B3994     177.571 144.536 123.116  1.00 20.00           P
TER
END
"""

def exercise_00():
  """ Verify that ligand is properly included by default """
  pdb_inp = iotbx.pdb.input(lines=pdb_str_1,source_info=None)
  hierarchy = pdb_inp.construct_hierarchy()
  asc = hierarchy.atom_selection_cache()
  ncs_inp = iotbx.ncs.input(hierarchy = hierarchy)
  ncs_groups = ncs_inp.get_ncs_restraints_group_list()
  assert len(ncs_groups) == 1
  assert len(ncs_groups[0].copies) == 1, len(ncs_groups[0].copies)
  # print list(ncs_groups[0].master_iselection)
  assert ncs_groups[0].master_iselection.size() == 58 # all chain
  # print hierarchy.select(asc.selection("chain A")).atoms_size()
  assert ncs_groups[0].copies[0].iselection.size() == 58 # all chain

if(__name__=='__main__'):
  exercise_00()
